python深度学习并行训练踩坑记录

最新推荐文章于 2025-10-20 16:38:08 发布

原创

最新推荐文章于 2025-10-20 16:38:08 发布 · 518 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #pytorch

本文记录了作者在实现深度学习并行训练过程中遇到的挑战，包括通过构造dataset实现多批量并行、多卡并行计算以及单卡多进程计算的方法。在尝试多批量并行时，遇到img和img_name匹配问题；多卡并行计算时，由于预训练模型参数限制无法直接应用；最后成功实现了单卡多进程计算，利用PyCharm的多进程功能实现并行。

最近在做一个对抗攻击的任务，写的代码只能逐个样本进行处理，速度实在太慢且显卡性能剩余的太多，于是想并行处理提高效率。

检索发现有以下实现方法：1.构造dataset实现多批量并行 2.多卡并行计算 3.单卡多进程计算。在尝试的过程中遇到了诸多问题，最终只实现了第三种方法。记录如下：

1.构造dataset实现多批量并行

根据https://www.cnblogs.com/Aiyuqianer/p/14091062.html理解了getitem函数的作用，编写代码如下，目的是传入一组图片及其文件名。

class ArcFaceDataset(Dataset):
    def __init__(self, samples_path,transformer=to_torch_tensor,loader=default_loader):
        self.sample_path=samples_path
        self.img_list = os.listdir(samples_path)
        self.transformer = to_torch_tensor
        self.loader=default_loader

    def __getitem__(self, idx):
        img_name=self.img_list[idx]
        img = self.loader(self.sample_path+'/'+img_name)
        img = self.transformer(img)
        return img,img_name

    def __len__(self):
        return len(self.img_list)

上面代码最终实现效果并不好，遇到的问题包括：传进来的img和img_name不匹配，img是一个四维的tenso