美文网首页
2020-10-29 Pytorch 程序单卡到多卡

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

作者: HellyCla | 来源:发表于2020-10-29 11:48 被阅读0次

    在多卡上训练的过程为先将模型和数据加载到第一张卡上,然后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,
                )
    

    相关文章

      网友评论

          本文标题:2020-10-29 Pytorch 程序单卡到多卡

          本文链接:https://www.haomeiwen.com/subject/orskvktx.html