2020-10-29 Pytorch 程序单卡到多卡

在多卡上训练的过程为先将模型和数据加载到第一张卡上，然后copy至其他卡。batchsize最好设为卡的整数倍，比如两张卡，bs为2， 那么每张卡分别计算bs=1的结果，在model.forward之后将不同卡上返回的结果合并传回再做下一步计算。

下面是踩坑后总结的需要注意的一些点。

1. 关于cuda与to(device)

• model 与 data 全部采用 to(device)方法来迁移至显存中，其中model.to(device)中的device, 默认是cuda:0就可以。 但是在网络计算过程中的中间数据，要放到哪个设备上是要根据当时在哪一块卡上计算决定的，因此，应将device设置为当前相关数据在的卡，x.to(related_data.device)。

2. 关于model与model.module

• 在调用model内部定义的变量或函数时，由于已经使用 model = nn.DataParallel(model)包裹，因此必须要写做model.module才能正常调用。 但是，调用model.forward()时必须要直接用model本身，否则无法使用多卡。

3. 关于多卡计算结果合并

• 必须确保forward最后return的均为tensor, 不能使用其他数据结构包裹，也不能返回标量，多个tensor分别返回。

4. 关于nn.ParameterList()不能用于多卡的问题

• 多卡不支持复制属性为List的参数，代码中定义的nn.ParameterList()会在多卡计算中无法复制到其他卡上，导致计算时参数为空的问题。需要写成module.register_parameter()的形式才能随module一起成功复制。示例如下

 for k in range(len(self.filters) + 1):
            # Weights
            H_init = np.log(np.expm1(1 / scale / filters[k + 1]))
            H_k = nn.Parameter(torch.ones((n_channels, filters[k + 1], filters[k])))  # apply softmax for non-negativity
            torch.nn.init.constant_(H_k, H_init)
            self.register_parameter('H_{}'.format(k), H_k)

            # Scale factors
            a_k = nn.Parameter(torch.zeros((n_channels, filters[k + 1], 1)))
            self.register_parameter('a_{}'.format(k), a_k)

            # Biases
            b_k = nn.Parameter(torch.zeros((n_channels, filters[k + 1], 1)))
            torch.nn.init.uniform_(b_k, -0.5, 0.5)
            self.register_parameter('b_{}'.format(k), b_k)

5. 其他

• 关于loss的计算部分不能放在model里，不然会产生不在同一张卡无法计算的问题
• 返回的tensor如果在原有网络中是按batch做了某种运算才返回的，这部分也需要加在model外另外做，因为多卡计算返回的是单独计算的结果。

更改为DistributedDataParallel

注意几个点：

• 模型和数据在两张卡上是一模一样的，相当于分别在两张卡上启动了两个python程序
  python -m torch.distributed.launch --nproc_per_node=2 脚本 其他命令行参数
• 一定要确保当前的模型和数据都设置了CUDA，在包裹为module之后一定还要加一句model.cuda()。全局可以写一个或者torch.cuda.set_device(local_rank)

            self.model=nn.parallel.DistributedDataParallel(
                self.model,
                device_ids=[opt.local_rank],
                output_device=opt.local_rank,
                broadcast_buffers=False,
            )