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

import torch

인공지능 모델은 학습을 할 때마다 성능이 달라진다.

분명히 이전 epoch에서는 loss도 낮고 train data에 대한 accuracy도 100% 였는데 다음 epoch에선 갑자기 loss가 높아지고 accuracy도 98%가 되는 경우가 있다.

이럴 때를 대비하여 학습 중간중간에 최적의 성능이라고 판단될 때 모델을 저장하는 것이 중요한데, 이때 torch 모듈에 있는 torch.save를 이용한다.

​

모델 저장하기

아래는 모델의 state_dict를 저장하는 방법으로 가장 권장되는 모델 저장 방법이다.

torch.save(model.state_dict(), PATH)

model = 모델 클래스명

'PATH' = 저장할 위치/저장할 이름.pth

​

모델을 저장할 때는 .pt 또는 .pth를 사용하는 것이 일반적이다.

eval을 위해 모델을 저장할 때는 학습된 모델의 매개변수만 저장하면 된다. 따라서 torch.save를 통해 model의 state_dict를 저장하는 것이 가장 권장된다. eval을 실행하기 전에는 반드시 model.eval()을 호출한다.

​

​

torch.save(model, PATH)로 저장할 경우, Python의 pickle 모듈을 이용하여 전체 모듈을 저장하게 된다.

그러나 이 때 pickle은 모델 그 자체를 저장하지 않기 때문에 사용한 특정 클래스/디렉토리 경로에 얽매이기 때문에 만들어둔 코드를 다른 프로젝트에서 사용하거나 리팩토링 한 후에 동작하지 않을 수 있다.

​

eval 또는 train 재개를 위해 checkpoint를 저장할 수도 있다. 이는 모델의 state_dict 보다 더 많은 것을 저장해야 한다. 자세한 내용은 이곳(https://tutorials.pytorch.kr/beginner/saving_loading_models.html)에서 확인할 수 있다.

​

모델 불러오기

model = ModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

어떤 모델인지 정의한 뒤, 저장한 모델을 불러오고 load_state_dict로 state_dict까지 불러오면 된다.

model eval 전에 model.eval() 을 선언해주는 것을 잊지 않도록 한다.