[pytorch] save pytorch model

 

Save

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

보통 torch.save를 사용해서 pytorch model을 저장하는데 이때 보통 .pt, .pth의 확장자를 쓴다. 그러나 .pth의 경우 python path와 충돌 위험이 있기때문에 .pt 확장자를 사용하는 것을 추천한다. 

Save for resume

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

일반적으로 모델을 저장할 때 학습 재개(resuming)를 위해 모델 파라미터 뿐만 아니라 optimizer, loss, epoch, scheduler등 다양한 정보를 함께 저장한다. 그래서 이러한 checkpoint는 모델만 저장할 때에 비해서 용량이 훨씬 커진다. 이럴때는 .tar 확장자를 사용한다. 나는 주로 .pth.tar를 사용한다.

To save multiple components, organize them in a dictionary and use torch.save() to serialize the dictionary. A common PyTorch convention is to save these checkpoints using the .tar file extension.

reference

• https://pytorch.org/tutorials/beginner/saving_loading_models.html
• https://stackoverflow.com/questions/59095824/what-is-the-difference-between-pt-pth-and-pwf-extentions-in-pytorch