python异步多线程 pytorch 多线程异步训练

转载

mob6454cc77b8eb 2023-08-30 23:40:55

文章标签 python异步多线程 pytorch 多GPU 多进程后向传播 文章分类 Python 后端开发

更正说明：

时间 2018/01/23

现在我所测试得到的dataParallel只能有效的进行前向传播，不能后向传播。下面的关于后向传播的原因我没有在代码层面进行分析，所以下面的记录部分（前面传播）是正确的。

不过可以得出的结论有：Dataparallel不支持自动的后向传播；Variable的grad_fn不能修改；Variable只有叶节点（在pytorch中计算图的倒着来的，叶节点就是初始输入节点）可以修改requires_grad属性，因此尝试从中间开始记录“计算图”的计划落空（希望不要引起误解，requires_grad属性并不是用来决定是否保留subgraph的，但有点联系；更何况网络中的参数也是叶节点，不单单是输入，比如图片）；

目前就是这样，搞清楚了再来记录下来。

pytorch多GPU训练

时间2018/01/22

pytorch支持多GPU训练，官方文档（pytorch 0.30）推荐使用DataParallel。但是这儿踩了一个很大的坑。

pytorch使用gpu的踩坑记录：

按照其官方网站上的说明，把神经网络结构做成了DataParallel对象，顺利的通过了前向传播，但是在后向传播进行梯度更新的时候出先了错误。后向计算以更新梯度时，因为grad_fn在dataparallel中，所以也要将更新梯度的方法放入到dataparallel中去。但是遗憾的是DataParallel计算得到的Variable中的属性grad_fn为torch.nn.parallel._function，而普通的网络所得到的对应grad_fn为torch.autograd.function。而前者不能用于backward()，所以最直接的想法是想修改计算图，但这样一点也不简单。

关于Variable和DataParallel：

简单的一提，这能帮助理解。Variable包含的的三个相关的项是data（记录数据Tensor），grad（记录梯度，以便后向传播），grad_fn（有人称为creator）。最后一个组成了计算所用的有向图。

DataParallel是torch.nn下的一个类，需要制定的参数是module（可以多gpu运行的类函数）和input（数据集），module需要有parameter方法，input数据需要是Variable类型，默认沿着0轴（参数dim，可修改）将input分成多份。

所以，pytorch（0.2.0）的多GPU训练方法：

将net和loss 以及nn.optim统统放入到新的自定义类中。但注意，DataParallel的参数module需要有paramer方法，看源码也看不明白parameter用来干嘛，有什么作用。但是踩坑过程中，继承nn.module的类和torch.nn.optim类都可以放入dataparallel中去，所以自定义类直接继承自nn.module就好，不管了。反正整个网络的前向后向传播都是在这个dataparallel中完成的。代码如下：

class train_parallel_module(nn.Module):#自定义类，继承自nn.Module，有人研究清楚了parameter欢迎讨论下啊
    """
    the net was used for DataParallel(module)
    """

    def __init__(self, net, optimizer, loss):#三个变量，见名知意
        super(train_parallel_module, self).__init__()
        self.net = net
        self.optim = optimizer
        self.loss = loss

    def forward(self, img, gt_heatmap):#img和gt_heatmap将沿着0轴分成等份。
        predicted = self.net(img)#前向传播
        l = self.loss(gt_heatmap, predicted)#loss计算

        # compute gradient and do SGD step
        self.optim.zero_grad()  # 后向传播
        l.backward()
        self.optim.step()  
        return l #结果返回，用于打印迭代的loss

调用时的代码：

# module for DataParallel
module = train_parallel_module(net, optimizer, loss)
module_parallel = nn.DataParallel(module, device_ids=device_ids)

imgs = Variable(imgs, requires_grad=True).cuda(device_ids[0])  # device_ids[0]
gt_heatmaps = Variable(gt_heatmaps).cuda(device_ids[0])
l = module_parallel(imgs, gt_heatmaps) # 在这儿调用
print 'loss is: {:.3f},and iter is {} with time {}'. \
    format(l.data[0], batch_iteration, time.time() - t)

既然说到了GPU，顺便说一下加载数据时的多进程吧：

data_loader = data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=6,
                                  pin_memory=True, drop_last=True, collate_fn=patch_data) #num_workers就是使用多进程