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

2020-10-29  本文已影响0人  HellyCla

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

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

1. 关于cuda与to(device)

2. 关于model与model.module

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

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

 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. 其他

更改为DistributedDataParallel

注意几个点:

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

猜你喜欢

热点阅读