Train with multi-gpu

GPU가 두개 이상인 경우, 모델 로드하고 학습/저장하는 방법 (부럽다)

command

# command
$ CUDA_AVAILABLE_DEVICES=0,1      (device numbers connected with comma)
# in __main__()
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"

model

Loading model

모델은 nn.DataParallel 이용해서 wrap해준다.
이때, 로드할 ckpt가 있다면 nn.DataParallel으로 wrap하기 전에 load해주어야 한다.

_model = MyModel()
if (args.ckpt is not None):
        checkpoint = torch.load(args.ckpt)
        _model.load_state_dict(checkpoint['model'], strict=False)
model = nn.DataParallel(_model).to(device)

위 코드 처럼 해야함 (모델 instance 생성 -> ckpt load -> dataparallel)
ckpt를 나중에 로딩했더니 그냥 깡통모델이 불러와졌다. 디버깅 하는데 고생 좀 했다 ㅜㅜ

Saving model

nn.DataParallel을 이용한 경우, model.state_dict()로 가중치를 불러올 수 없다. model.module.state_dict()로 접근해야 한다고 함.
그래서 모델 저장 시,

torch.save({'model': model.module.state_dict()}}

로 해야 함.
multi gpu이든 아니든 .module로 접근하는게 더 general하다고 한다.

https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-torch-nn-dataparallel-models

더 많은 gpu, dataloader worker를 사용한다고 무조건 좋은 것은 아니다. cpu에서 처리하는 연산량을 고려하지 않으면 병목현상이 발생할 수 있음. 짬날때 관련해서 공부좀 해야겠다.