Pytorch DistributedDataParallel Imagenet trainning

pytorch imagenet multiprocessing-distributed training

 

imagenet을 학습 시킬 일이 있어서 작성....

 

보통 사용하는 torch.nn.DataParallel을 사용할 시 GPU-Util을 확인해 보면, 50% 정도밖에 사용하지 못하는 모습을 볼 수 있다.

DistributedDataParallel를 사용하면 이를 해결 할 수있고, 좀더 빠르게 학습이 가능하다....

원본 코드는 아래의 pytorch imagenet official 이며, 이를 바탕으로 수정

https://github.com/pytorch/examples/tree/master/imagenet

 

원본 코드에선 사용하는 GPU 수 만큼 log들이 모두 출력되도록 작성되어 있다.

몇몇 코드를 수정 하고 추가해서 이를 해결 했다.

 

기존 코드 log

수정 코드 log

실행 commander : -save_path './test/' --gpu_count 4 이 부분만 수정하고 돌리면 된다.

-save_path : 저장 경로, --gpu_count : 사용할 gpu 갯수이다.

 

python main.py -a resnet50 --dist-url 'tcp://127.0.0.1:2222' --dist-backend 'nccl' --multiprocessing-distributed --world-size 1 --rank 0 none --save_path './test/' --gpu_count 4

 

reduce_tensor 에서 사용하는 gpu log들을 모아 평균 계산.

 

전체 코드

https://github.com/MOONJOOYOUNG/pytorch_imagenet_multiprocessing-distributed

for i, (images, target) in enumerate(val_loader):
  data_time.update(time.time() - end)

  if args.gpu is not None:
  images = images.cuda(args.gpu, non_blocking=True)
  target = target.cuda(args.gpu, non_blocking=True)

  output = model(images)
  loss = criterion(output, target)

  acc1, acc5, correct = util.accuracy(output, target, topk=(1, 5))

  reduced_loss = reduce_tensor(loss.data)
  reduced_top1 = reduce_tensor(acc1[0].data)
  reduced_top5 = reduce_tensor(acc5[0].data)

  losses.update(reduced_loss.item(), images.size(0))
  top1.update(reduced_top1.item(), images.size(0))
  top5.update(reduced_top5.item(), images.size(0))
	
  if dist.get_rank() == 0:
    if i % args.print_freq == 0:
      progress.display(i)

def reduce_tensor(tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.reduce_op.SUM)
    # gpu 갯수로 나눠줌.
    rt /= args.gpu_count
    return rt