如何并行运行程序

参考了官方文档， torch.nn.parallel.DataParallel
以及https://zhuanlan.zhihu.com/p/102697821

在运行此DataParallel模块之前，并行化模块必须在device_ids [0]上具有其参数和缓冲区。在执行DataParallel之前，会首先把其模型的参数放在device_ids[0]上，一看好像也没有什么毛病，其实有个小坑。我举个例子，服务器是八卡的服务器，刚好前面序号是0的卡被别人占用着，于是你只能用其他的卡来，比如你用2和3号卡，如果你直接指定device_ids=[2, 3]的话会出现模型初始化错误，类似于module没有复制到在device_ids[0]上去。

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

当你添加这两行代码后，那么device_ids[0]默认的就是第2号卡，你的模型也会初始化在第2号卡上了，而不会占用第0号卡了。这里简单说一下设置上面两行代码后，那么对这个程序而言可见的只有2和3号卡，和其他的卡没有关系，这是物理上的号卡，逻辑上来说其实是对应0和1号卡，即device_ids[0]对应的就是第2号卡，device_ids[1]对应的就是第3号卡。（当然你要保证上面这两行代码需要定义在

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

而且不要以为添加了to(device)或者.cuda（） 就运行在卡0上了，其实不然！ 这只是将模型暂时放到了卡0上，第二步系统会复制模型和参数的！

而且并行和to、cuda是缺一不可的！

代码测试：

import torch
import torch.nn as nn
import torch.nn.functional as F
import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  #指定程序可见的devices！ 也就是其它设备都不可见
os.environ["CUDA_VISIBLE_DEVICE"] = "0, 1"    #设置当前使用的GPU设备为1,0号两个设备,名称依次为'/gpu:0'、'/gpu:1'。表示优先使用1号设备,然后使用0号设备

## 看起来没事，但是有坑的！ 如果我们和别人一块使用的话，别人在用卡0，因为程序会默认将模型参数和缓存放到卡0，那么你不改上面的话就会出现问题

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))


input = torch.randn(128, 1, 256, 128)
net = Model()
if torch.cuda.device_count() > 1:
    # model.to(device)
    device_ids = [0, 1, 2]
    net = torch.nn.parallel.DataParallel(net, device_ids=device_ids)
    net.cuda()
    # net.to(device)    //两种方法都可以！ 只不是这种方法考虑了不能用的情况！
    for param in next(net.parameters()):
        print(param, param.device)
    print(next(net.parameters()).device)
net = net(input)
print(net)