[Pytorch] AttributeError: 'DataParallel' object has no attribute 'predict'

파이썬에서 딥러닝 프레임워크를 pytorch로 쓰다가 GPU 병렬처리를 위해서 DataParallel 처리를 해주었다.

 

import torch

# model 생성부분 생략
model = torch.nn.DataParallel(model, device_ids=[0,1,2,3]) # GPU 0,1,2,3 총 4개 사용
model.cuda()

 

모델을 학습까지 완료시킨 후 다음과 같이 model.predict()을 하니 아래와 같은 에러 메세지가 나왔다.

 

코드

# x_tensor 생성부분 생략
pr_mask = model.predict(x_tensor)

에러 메시지

AttributeError: 'DataParallel' object has no attribute 'predict'

 

해결 방안은 다음과 같이 model과 predict 사이에 .module을 추가하여 코드를 수정하면 된다.

pr_mask = model.module.predict(x_tensor)

 

이상없이 잘 돌아간다!

 

 

참고한 github issue

https://github.com/jytime/Mask_RCNN_Pytorch/issues/2

AttributeError: 'DataParallel' object has no attribute 'train_model' · Issue #2 · jytime/Mask_RCNN_Pytorch

 

 

 

 

-히비스서커스-