[Pytorch] 모델 저장하고 불러오기

파이토치 한국사용자 모임의 글을 보면서 공부한 내용을 정리한 글입니다.

 

 

모델을 저장하고 불러오기를 통해 모델의 상태를 유지할 수 있다.

 

모델 가중치 저장

Pytorch 모델은 학습한 매개변수를 state_dict라고 불리는 내부 상태 사전(Internal State Dictionary)에 저장한다.

이 상태 값들은 torch.save 메소드를 사용하여 저장할 수 있다.

import torch
import torchvision.models as models

model = models.vgg16(weights='IMAGENET1K_V1')
# torch.save(모델의 상태 사전, 저장할 경로) 
torch.save(model.state_dict(), 'model_weights.pth')

 

 

모델 가중치 불러오기

먼저 동일한 모델의 인스턴스를 생성한 다음 load_state_dict() 메소드를 사용하여 모델의 가중치를 불러온다.

model = models.vgg16() # 여기서는 ``weights`` 를 지정하지 않았으므로, 학습되지 않은 모델을 생성합니다.
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()
"""
추론(inference)을 하기 전에 model.eval() 
메소드를 호출하여 드롭아웃(dropout)과 배치 정규화(batch normalization)를 
평가 모드(evaluation mode)로 설정해야 합니다. 그렇지 않으면 일관성 없는 추론 결과가 생성됩니다.
"""

 

모델 자체를 저장 하고 불러오기

저장 : .save 메소드에 모델 인스턴스를 전달하고 

불러오기 : .load(저장한 파일 이름)

• 이 방식은 Python의 pickle 모듈을 사용하여 모델을 직렬화하므로, 모델을 불러올 때 실제 클래스 정의를 적용한다.

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

 

7. Model Save&Load.ipynb