Pytorch 5 : Save and Load

import torch
import torchvision.models as models

 

이번 글은 학습된 모델을 저장하고, 불러오는 방식에 대한 코드이다.

 

파이토치 모델은 학습된 파라미터를 가지고 있고 이 파라미터들을 파이썬 자료형인 Dictionary 형태로 저장하고 있다.

이를 우리는 'state_dict'이라 한다.

이를 우리는 'torch.save()'를 사용하여 저장할 수 있다.

 

model = models.vgg16(weights = 'IMAGENET1K_V1')
torch.save(model.state_dict(), 'model_wights.pth')

 

model 변수는 ImageNet을 활용하여 사전학습된 모델 VGG16을 가리키고 있다.

물론 사용자가 자신이 직접 학습시킨 모델을 활용하고 싶다면 이 부분을 바꿀 수 있다.

이 예시에서는 torchvision에 저장된 모델을 가져오는 코드를 보여주고 있다.

이 모델을 우리 로컬 환경에 저장하고 싶다면, torch.save()를 활용해서 state_dict 형태로 학습된 파라미터를 저장할 수 있고 파일명은 'model_weights.pth'이다.

 

그러면 이렇게 로컬 환경에 저장된 모델을 다시 불러오고 싶다면 어떻게 해야 할까?

'load_state_dict()'을 사용하면 된다.

 

model = models.vgg16()
model.load_state_dict(torch.load('model_weights.pth'))

 

위의 model 변수명과 아래의 model 변수명의 차이는 무엇일까?

weights를 따로 명시적으로 지정해주지 않았다. 이러면 디폴트로 랜덤하게 가중치가 생성된 VGG16 모델이 불러와진다.

이 모델의 파라미터를 우리가 앞서 저장한 파라미터 'model_weights.pth'로 변경하고 싶다면,

위와 같이'load_state_dict()'을 사용해주면 된다.

그러면 model의 파라미터가 앞서 ImageNet으로 학습된 VGG16의 파라미터로 설정된다.

 

이처럼 학습된 파라미터를 특정한 모델에 부여하고 싶다면, 우리는 먼저 동일한 모델 (동일한 파라미터 구조를 가진 모델) instance를 생성해야만 한다. 그 뒤, 해당 모델에 학습된 파라미터를 부여하는 방식으로 사용하면 된다.

 

torch.save(model, 'model.pth')
model = torch.load('model.pth')

 

맨 위 코드 'torch.save(model.state_dict())'과 바로 위 코드 'torch.save(model)'의 차이점은 무엇일까?

모델의 학습된 파라미터뿐만이 아니라 모델의 구조 자체도 함께 저장하고 싶을 수 있다.

모델 파라미터를 저장하기 위해서는 우리는 먼저 모델 클래스를 instance로 생성한 후, (그러면 이 instance에는 모델 구조도 함께 정의되어져 있다.) state_dict() 메소드를 사용하는데,  이때 모델의 구조를 함께 저장하고 싶다면 model 인스턴스 자체를 변수에 집어넣으면 된다.

 

즉, model.state_dict()이 아닌 model 자체가 save() 메소드의 변수로 들어가면 모델의 학습된 파라미터뿐만이 아니라 구조 자체도 함께 저장을 할 수 있다.

 

이것보다 조금 더 일반적인 방식의 저장 방법은 무엇이 있을까?

우리는 학습 도중 checkpoint를 설정하여 그 지점에서 모델을 저장하거나 불러오길 원할 수도 있다.

만약 학습을 도중에 중단하였다가 다시 이어서 진행하거나, 특정한 에폭 또는 loss 값을 기준으로 모델을 저장하고 싶을 수도 있다. 혹은 학습 도중의 한 지점 (epoch)에서의 파라미터뿐만이 아니라 여러 지점(epoch)에서의 파라미터들을 다량으로 저장해두고 싶을 수도 있다.

이를 위해서는 조금 더 디테일한 코드 구현이 필요하다.

 

 

우선, 우리가 원하는 저장 지점들을 Dictionary 형태로 선언해준다. 파이토치는 이 지점에서 모델을 .tar 파일로 저장해두는데 이것을 사용하기 위해서는 위와 동일하게 모델을 선언해줄 뿐만이 아니라, optimizer도 선언해주어야 한다.

 

class Net(nn.Module):
	def __init__(self):
    	super(Net, self).__init__()
        
	def forward(self, x):
    	return x

net = Net()

 

위와 같이 정의된 모델이 있다고 해보자.

우선 __init__() 생성자 메소드에 layer가 정의돼 있을 것이고 forward() 메소드에 연산 과정이 정의될 것이다.

이는 이전 글에서 다루었으므로 skip하자.

 

다음으로 SGD with Momentum optimizer를 선언해준다.

 

import torch.optim as optim

optimizer = optim.SGD(net.parameters(), lr = 0.001, momentum = 0.9)

 

다음으로 우리가 저장하고 싶은 상황과 관련된 정보를 Dictionary 형태로 선언해준다.

 

epoch = 5
path = 'model.pt'
loss = 0.1

torch.save({'epoch' : epoch, 'model_state_dict':net.state_dict(), 'optimizer_state_dict':optimizer_state_dict(), 'loss':loss,}, path)

 

이렇게 저장된 모델을 불러올때는 다음과 같이 해주면 된다.

model = Net()
optimizer = optim.SGD(model.parameters(), lr = 0.001, momentum = 0.9)

checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']


model.eval()
# or
model.train()

 

먼저 model과 optimizer를 선언해준다음에 저장된 모델을 불러와주어야 한다.

만약 학습을 다시 이어서 진행하고 싶다면, model.train()을 사용해주고 만약 불러와진 모델을 활용해서 평가를 하고 싶다면 model.eval()을 사용해주면 된다.