在查看SlowFastNet源代碼https://github.com/facebookresearch/SlowFast中的model_builder時,想要採用多GPU訓練;GPU大於1的話,代碼自動調用torch.nn.parallel.DistributedDataParallel套在model外面。
torch.nn.parallel.DistributedDataParallel和torch.nn.parallel.DataParallel作用相似,都是pytorch調用多GPU訓練;但torch.nn.parallel.DistributedDataParallel需有個初始化
在torch.nn.parallel.DistributedDataParallel前加入
torch.distributed.init_process_group('nccl',init_method='file:///home/.../my_file',world_size=1,rank=0)
這裏是在單個機器上調用多張GPU,簡稱單機多卡,所以world_size=1;具體參考
https://github.com/pytorch/examples/tree/master/imagenet
下列代碼構造了僅使用index爲1,2共計兩個GPU同時訓練模型
os.environ['CUDA_VISIBLE_DEVICES']='1,2'
device=torch.device('cuda:0')
model=ConvNet(num_classes)
torch.distributed.init_process_group('nccl',init_method='file:///home/.../my_file',world_size=1,rank=0)
model=torch.nn.parallel.DistributedDataParallel(model.to(device)
關於pytorch多GPU訓練,下面這篇文章寫得有點亂,看了下就是把他放在最後的幾篇文章給結合了一下
https://blog.csdn.net/m0_38008956/article/details/86559432