20201114-pytorch指定GPU

使用pytorch的并行GPU接口

model= torch.nn.DataParallel(model, device_ids=[0])

指定多卡：
model = torch.nn.DataParallel(model, device_ids=[0,1,2,3])

先看看我的卡：

图片.png

可以看到我有六个，所以：
model = torch.nn.DataParallel(model, device_ids=[0,1,2,3,4,5])
但是我发现有两个竟然没用上，后来发现是batch_size太小，所以调整之后就跑满了：

只要这样一个简单的包裹，Pytorch已经为我们做了很多复杂的工作。我们只需要增大我们训练的batch_size(一般计算为N倍，N为显卡数量)，其他代码不需要任何改动,我的是6块卡，所以设成20，就只能用5块，所以得是24

图片.png

这里需要说明的是，device定义中的cuda:0可以不变，无论你多卡的GPU如何选取。device_ids=[0,1,2,3]也相对固定，如果有两张卡就定义device_ids=[0,1]，如果有八张卡就定义device_ids=[0,1,2,3,4,5,6,7]。在多卡并行时，会有一个相对的主卡，就是cuda:0所指向的GPU。介绍到主卡就得引入os.environ[“CUDA_VISIBLE_DEVICES”] = “0,1,2,3”，在这里主卡是0卡。如果定义为cuda:1，os.environ[“CUDA_VISIBLE_DEVICES”] = “1,2,3”，这里主卡是2卡。
提示：需要注意的是device_ids中包含的卡数要和os.environ[“CUDA_VISIBLE_DEVICES”]中包含的卡数相等。这里的卡数是指数量，无需具体卡号一一对应，此外，batch_size的数量需要大于GPU的数量。

OK，最核心的部分给大家介绍完了，剩下的工作就是将之前单卡运行时的所有.cuda()替换为.to(device)即可。

需要使用.cuda或者.to(device)的地方。

model（如：model.to(device)）