Pytorch 模型训练如何提速（speed up pytorch model training）

最新推荐文章于 2025-06-21 15:03:42 发布

原创最新推荐文章于 2025-06-21 15:03:42 发布 · 3.5k 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#Pytorch #训练加速 #speed up model training #数据导入加速 #Pytorch模型训练

深度学习专栏收录该内容

37 篇文章

订阅专栏

本文探讨了PyTorch模型训练速度慢的三大原因，并提供了详细的优化策略，包括数据导入优化、模型及loss函数优化技巧，以及使用半精度训练和多GPU等高级方法。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前言

导致pytorch的模型训练速度比较慢的原因最有可能的是三个：1. 数据导入环节，操作复杂 2.模型本身很复杂，数据流在模型中传递时过于耗时 3.loss函数计算复杂。

这其中第一个环节往往是最有可能的原因，第二，三个环节其实一回事；pytorch本身的框架针对这两个问题也做了大量的优化，如果不是专业技术过硬，建议在这两个环节上就不要过于纠结了，设计简单易用的模型才是正道。

数据导入环节的优化

数据导入环节，尤其是诸如图像等大张量从内存中的反复读取，以及后续的数据增广操作往往是造成训练速度低的主要原因。针对这个环节的加速其实有一些trick可以试用。

常规的做法

常规的做法主要是训练中各种现成的pytorch工具使用以及训练参数的设置，主要有如下的几种方案：

采用pytorch 自带的Dataloader而不是自己编写张量的导入类，Dataloader可以方便的设置cpu多线程，很多操作诸如张量的缩放都经过了优化
采用更大的batch size，设置cpu或者GPU能够承受的最大batch size，这最大的用处在于节省了后续梯度传递时花的时间
可以使用累积梯度，其实就是在cpu或者GPU能承受范围内，多次循环batch 再进行梯度计算
保存图，就像下面这样

losses = []
...
losses.append(loss)
print(f'current loss: {torch.mean(losses)'})

5. 使用多个GPU

当然这里推荐一个用于加速运行的插件，叫做pytorch lighting （https://github.com/williamFalcon/pytorch-lightning）

它的使用也是比较简单的

from pytorch_lightning import Trainer
model = LightningModule(…)
trainer = Trainer()
trainer.fit(model)

非常规的做法

非常规的做法就得视场合而定了，这些做法并不是对所有的应用场景有效，在不适合的场景里可能造成严重的训练质量下降。

半精度或者混合精度训练，该方法在一些本来对张量精度要求不是很高的领域比较适用，可以显著的提高训练速度，同时显著降低运算显存开销，但是并不是所有领域都适合。关于半精度以及混合精度，可以采用apex library 在英伟达显卡上方便的实现。
提前规范化数据，比如大量的图像张量导入时，可以将图像提前缩放成2的n次方类型，这主要是因为大量的框架优化对于这种尺度的图像处理优化效果明显，而对于任意尺寸输入的图像不敢保证；但是提前的规范化有可能造成一些细节的变形
使用hdf5格式，提前将数据转成hdf5格式，这种格式对于cpu运算较为友好，同时也是受限比较小的一种方式；但是我在使用中发现，hdf5的解析有赖于自己写的方式，如果技术不过硬这里有可能还是解决不了问题。我这里有一个示例类，可供参考

class AdobePatchDataHDF5(data.Dataset):

def __init__(self, root, cropsize = 256, outputsize = 256):

    fgfile = h5py.File(root, 'r')

    self.root = root
    self.fgfile = fgfile
    self.cropsize = cropsize
    self.outputsize = outputsize


def __getitem__(self, index):
    # read image
    fgimg = self.fgfile['img'][index, ...]

    # random crop and resize, random flip with cv2

    # toTensors
    fgimg = fgimg.astype(np.float32) / 255.0 
    fgimg = torch.from_numpy(fg.transpose((2, 0, 1)))

    # norm [0, 1] to [-1, 1]

    return fgimg, label

def __len__(self):
    return self.fgfile['img'].shape[0]

还有一些是同gpu绑定的方法，比如使用Nvidia DALI，（https://github.com/NVIDIA/DALI），这在预处理阶段可以进行极大的加速，但是目前的稳定版本(截止12.25)好像只能支持有限型号的显卡。