[Pytorch] model save 후 load (feat. multi GPU 사용 저장 후 single GPU로 load 하여 evaluation하는 방법)

torch framework로 model을 train 후 저장 및 불러오는 방법은 크게 두 가지가 있다.

 

1) torch.save(model.state_dict(), model_path)

 

이 방법은 모델의 가중치만 저장하는 방법이다. 

 

모델 train 후 evaluation 모드 전환 후 가중치만 저장

import torch.nn    as nn
import torchvision
import torch

# 예시 모델
model = torchvision.models.resnet18()

# GPU 2대를 사용한다고 가정 후 dataparallel을 통한 병렬처리
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model = nn.DataParallel(model, device_ids=[0,1])

### train 과정 생략 ###

# 모델을 eval 모드로 전환
model.eval()

# 가중치 저장
torch.save(model.state_dict(), './model_path.pth')

다른 스크립트에서 불러와 evaluation

import torch.nn      as nn
import torchvision
import torch

# 불러올 가중치의 모델을 다시 생성
best_model = torchvision.models.resnet18()

# 다시 생성한 모델에도 동일한 이전의 모델과 동일한 처리
device = "cuda" if torch.cuda.is_available() else "cpu"
best_model.to(device)
best_model = nn.DataParallel(best_model, device_ids=[0,1])

# 생성한 모델에 저장한 가중치를 불러옴
best_model.load_state_dict(torch.load('./model_path.pth'))

### eval 과정 생략 ###

 

만약, 생성한 모델에 이전의 모델과 동일한 처리를 해주지 않고 (GPU처리나 dataparallel 등) 가중치를 불러온다면 에러가 발생한다. 따라서, 학습은 multi GPU를 사용하고 evaluation할 시에는 single GPU를 사용하고 싶은 경우 적합하지 않기에 두번째 방법으로 저장 후 불러와 해결할 수 있다.

 

 

2) torch.save(model, model_path)

 

이 방법은 모델과 가중치를 둘 다 저장하는 방법이다.

 

모델 train 후 evaluation 모드 전환 후 모델과 가중치 저장

import torch.nn    as nn
import torchvision
import torch

# 예시 모델
model = torchvision.models.resnet18()

# GPU 2대를 사용한다고 가정 후 dataparallel을 통한 병렬처리
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model = nn.DataParallel(model, device_ids=[0,1])

### train 과정 생략 ###

# 모델을 eval 모드로 전환
model.eval()

# 모델과 가중치 저장
torch.save(model, './model_path.pth')

다른 스크립트에서 불러와 evaluation

import torch.nn      as nn
import torchvision
import torch

# 저장한 모델과 가중치를 불러옴
best_model = torch.load('./model_path.pth')

# GPU만 처리
device = "cuda" if torch.cuda.is_available() else "cpu"
best_model.to(device)

### eval 과정 생략 ###

이 방법을 사용하면 학습은 multi GPU를 사용하고 evaluation할 시에는 single GPU를 사용하는 처리를 할 수 있다. 만약, 모델을 학습 후 첫 번째 방법으로 저장하였더라도 첫 번째 방법으로 모델을 불러온 후 두 번째 방법으로 다시 저장하여 불러오면 처리가 가능하다.

 

그렇다면 두 번째 방법이 훨씬 더 간단하고 자유로운 것이 아니냐 할 수 있지만 두 번째 방법 사용 시에는 주의할 점이 있다. model을 직접 코드로 작성한 모델을 train 후 저장하여 다른 경로나 서버에서 불러올 경우이다. 이때에는 저장할 때 학습스크립트에서의 모델 스크립트 경로가 불러오는 스크립트 경로에도 동일하게 있어야 에러가 발생하지 않는다는 것이다.

 

예시를 들어보자면

 

project 

     ㅏtrain

     ㅣ ㅏtrain.py

     ㅣ  Lmodel/resnet.py

      L eval

         ㅏeval.py

          Lmodel/resnet.py

 

train.py에서 학습 후 모델 저장 후 eval.py에서 모델을 불러올 시 eval directory의 model/resnet.py가 없다면 (혹은 train directory의 경로와 다르다면) 에러가 발생하므로 항상 유의해주어야 한다. 다른 방법도 존재하는데 

 

https://biology-statistics-programming.tistory.com/213?category=973653 

[error] torch.load error (feat. ModuleNotFoundError: No module named 'models')

 

위의 글에 간단하게 정리해두었다.

 

 

 

- 히비스서커스 -