How to use a single-machine multi-GPU to run

[CoderwayNew]

The function to get the device: init_distributed_device(args), in training.main.py line 129, it seems that only a single-GPU device can be obtained. The key part of the function is defined as follows:

  if torch.cuda.is_available():  
        if args.distributed and not args.no_set_device_rank: 
            device = 'cuda:%d' % args.local_rank 
        else:  
            device = 'cuda:0'  
        torch.cuda.set_device(device)  
    else:
        device = 'cpu'

How to make the project run on a single-machine multi-GPU?

[CoderwayNew]

Just use the command line, then enter the project directory:

torchrun --nnode=1 --nproc_per_node=2 ./training/main.py

--nnode=1 # single-machine --npor_per_node =2 # 2 GPU