accelerator.prepare
accelerate의 가이드를 보고 현재 학습 중이던 모델을 multi gpu를 통해 학습하기 위해 코드를 아래와 같이 작성했다.
model, optimizer, train_loader, criterion = accelerator.prepare(model, optimizer, train_loader, criterion)Issue
문제가 발생하지 않을 것으로 생각했으나,
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: loss += torch.mean(self.match_loss(pred, gt))
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: gt_expand = torch.gather(gt, 1, feature_id).view(batch_size, self.pnum, self.pnum, 2)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:2! (when checking argument for argument index in method wrapper_CUDA_gather)위와 같은 에러가 지속됐다. loss 계산 중 tensor들이 같은 device에 있지 않은 문제였다. criterion을 accelerate.prepare로 넘겨주었기 때문에, 자동으로 criterion 내부의 데이터들이 device 매칭 후 계산될 것이라 생각했는데, 아니었다.
accelerate.prepare은 모델에 관련된 파라미터와 옵티마이저를 이동시키는 역할을 한다. 문제가 생긴 Loss는 직접 class 인스턴트를 만들었기 때문에, accelerate.prepare에서 device로 옮기지 않아서 생긴 문제였다.
self.feature_id = self.feature_id.to(accelerator.device)
gt = gt.to(self.accelerator.device)이렇게 loss class에 accelerator를 전달하고, 이 accelerator의 device로 데이터를 옮겨 해결할 수 있었다.
또한, 저장 과정도 일부 수정해주어야 한다.
if accelerator.state.distributed_type == "MULTI_GPU":
model = accelerator.unwrap_model(model)모델을 저장하기전, 위와 같이 accelerator를 통해 unwrap을 해주고 나서 저장해야 기존과 동일하게 load 할 수 있다.
