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()을 사용해주면 된다.
'Deep dive into Pytorch' 카테고리의 다른 글
Pytorch 7 : Tensor 심화 (0) | 2023.07.23 |
---|---|
Pytorch 6 : Implement CNN (0) | 2023.07.20 |
Pytorch 4 : Training (1) | 2023.07.16 |
Pytorch 3 : Neural network 구현 (0) | 2023.07.15 |
Pytorch 2 (0) | 2023.07.11 |
댓글