학습은 cuda:2 환경에서 진행한 모델을 다른 pc에서 불러와 cuda:0 환경에서 inference를 시도
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = Net()
net.to(device)
net.load_state_dict(torch.load("./output/trained_model.pth"))
이러한 경우 torch.load 에서 에러가 발생.
학습 당시의 환경은 cuda:2 인데 그 환경과 다르다는것!
이럴 경우 map_location 파라미터를 설정 해 주면 정상적으로 읽기 가능
net.load_state_dict(torch.load("./output/0074.pth", map_location='cuda:0'))
'전공관련 > Deep Learning' 카테고리의 다른 글
| [Pytorch] Multi GPU를 활용 해 보자 (2) | 2019.03.15 |
|---|---|
| [Pytorch] network 결과물 사용시 주의사항! (0) | 2019.03.13 |
| [Pytorch] pretrained-model 의 일부만을 불러와 보자 (0) | 2019.03.12 |
| [Pytorch] 학습 한 모델을 저장하고 불러오자 (1) | 2019.03.12 |
| Deconvolution 파라미터에 따른 출력 크기 계산하기 (0) | 2018.10.19 |
매직블럭
작은 지식들 그리고 기억 한조각