[파이토치/Pytorch] torch.save(), checkpoints, 전이학습(Transfer Learning) 방법 기초

1) torch.save()

학습시킨 모델을 저장하고 불러오는 것은 빈번하게 사용하는 기능이다. 아래의 두 방법이 있다.

 

① torch.save(model.state_dict(), 경로)

# 모델의 파라미터만을 저장
torch.save(model.state_dict(), os.path.join(MODEL_PATH, "model.pt"))

# 같은 구조의 모델에 저장했던 파라미터 load
new_model = TheModelClass()
new_model.load_state_dict(torch.load(os.path.join(MODEL_PATH, "model.pt")))

모델의 구조가 모듈 형태로 이미 존재하거나 별도로 저장해놓는 경우가 많기 때문에 보통 위 방법을 사용하게 된다.

 

② torch.save(모델, 경로)

# 모델의 아키텍쳐까지 저장 -> 더 큰 용량을 필요로 함
torch.save(model, os.path.join(MODEL_PATH, "model.pt"))

# 저장했던 모델과 파라미터 load
model = torch.load(os.path.join(MODEL_PATH, "model.pt"))

 

모델의 구조까지 한 번에 저장하고 불러올 수도 있다.

 

2) Checkpoints

모델 학습은 보통 긴 시간 이루어지고, 그 중간에도 학습이 올바른 방향으로 진행되고 있는지를 확인하거나, early stopping을 하거나, 모델 성능이 좋아지는 경우에만 저장하는 등의 목적을 위해 지속적으로 학습 과정을 모니터링하는 것이 필요하다.

# 모델의 정보를 epoch과 함께 저장, 파일 명에 예시처럼 핵심 정보들을 적어주는 것이 일반적이다.
torch.save({'epoch': e,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': epoch_loss,
},
f"saved/checkpoint_model_{e}_{epoch_loss/len(dataloader)}_{epoch_acc/len(dataloader)}.pt")

# 저장한 checkpoint 불러와 확인
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']

마찬가지로 torch.save()를 활용한다. epoch과 loss, parameter와 optimizer의 정보까지 함께 저장하는 것이 보통이며, 그 파일명에서 핵심적인 내용을 파악할 수 있도록 하는 것이 좋다.

 

3) Transfer learning(전이학습)을 위한 gradient freeze

이전에 다른 데이터셋으로 학습한 pre-trained 모델을 불러와 내 데이터를 추가로 학습시키거나 내 데이터 상황에 맞게 일부를 바꾸어 학습하는 것을 transfer learning(전이학습)이라 한다. 이를 위해 원래 모델의 앞, 뒤에 필요한 layer를 추가하고 gradient들을 필요에 따라 freeze하는 방법을 알아야한다.

 

# vgg19 모델 불러오기
vgg = models.vgg19(pretrained=True).to(device)

# 모델의 마지막 부분에 원하는 목적에 맞게 layer 추가
class MyNewNet(nn.Module):
    def __init__(self):
        super(MyNewNet, self).__init__():
			self.vgg19 = models.vgg19(pretrained=True)
			self.linear_layers = nn.Linear(1000, 1)

		# Defining the forward pass
		def forward(self, x):
			x = self.vgg19(x)
			return self.linear_layers(x)

# 마지막에 추가한 layer를 제외하고 나머지 모든 파라미터에 대해 requires_grad = False를 선언하므로써 freeze한다.
for param in my_model.parameters():
		param.requires_grad = False
for param in my_model.linear_layers.parameters():
		param.requires_grad = True

위 예시는 vgg19를 불러와 활용하되, 둘 중 하나로 분류하는 목적에 맞게 맨 마지막에 출력 노드가 1인 layer를 추가한 예시 코드이다. (추후에 sigmoid 함수를 통과시켜 분류 모델을 완성하게 된다.) layer를 init에 추가해주고 나면 forward 함수도 수정해줘야한다.

(이렇게 레이어를 추가하거나, 경우에 따라 마지막 fc layer를 분석하여 그 자체를 변경하는 방법도 있다.)

 

추가로 마지막 두 번의 for문을 통해 본래 vgg19에 있던 파라미터들은 freeze하고, 내가 추가한 layer의 파라미터들만 requires_grad = True로 설정하여 backward 연산에 의해 변하도록 설정한다.