[PyTorch] 모델 저장하기 & 불러오기

state_dict

- torch.nn.Module 모델의 각 계층의 학습 가능한 매개변수(model.parameters())들을 매핑하는 dictionary 객체.

- torch.optim 옵티마이저 객체 또한 옵티마이저의 상태와 사용된 하이퍼파라미터 정보를 포함한 state_dict를 가짐.

 

---

 

state_dict 저장하기/불러오기

저장하기

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

불러오기

model = TheModelClass(*args, **kwargs)
model.load_sate_dict(torch.load(PATH)
model.eval()

- inference를 위해 학습된 모델의 학습된 매개변수만 state_dict를 이용하여 저장하는 방법.

- 모델 저장 시 .pt 또는 .pth 확장자를 사용하는 것이 일반적인 규칙임.

- inference 실행 시 반드시 model.eval()을 호출하여 평가 모드로 설정하여야 함.

 

---

 

전체 모델 저장하기/불러오기

저장하기

torch.save(model, PATH)

불러오기

model = torch.load(PATH)
model.eval()

- 전체 모델을 저장하고 불러오는 방법.

- 모델 그 자체를 저장하지 않기 때문에 직렬화된 데이터가 모델을 저장할 때 사용한 특정 클래스 및 디렉토리 경로에 얽매인다는 단점이 있음.

- 모델 저장 시 .pt 또는 .pth 확장자를 사용하는 것이 일반적인 규칙임.