1 模型训练时出现的问题

内存容量和磁盘IO速度影响GPU利用率，进而影响模型训练速度

如果要想在模型训练时利用GPU，提高GPU的利用率，不受限于IO瓶颈，一般常用的方式是：

  • 数据预加载，将所有数据都预加载到内存中，内存的读写速度很快，从内存传递数据到GPU速度也很快
  • 数据预先离线处理，然后存储在固态硬盘中，在模型训练时，从固态硬盘中读取数据传递到GPU中，但是其实这种方式也是受限于固态硬盘的读取速度

那么先描述下我今天面对的问题，我最近在做一个图片的二分类问题，自己建了一个100W+的图片大数据集，每张图片分辨率为256x256。

1.1 数据预加载到内存中-受限于内存容量

如果我们采取数据预加载的方式，需要将100W+的图片预先加载到内存中，那么100W+图片需要多少内存呢？

我们这里可以算一下，每一张图片占用内存256x256x3 = 196608个字节，那么100W图片需要1000000x196608=196608000000个字节=192000000kb=187500mb=183.1gb，但是我本地的内存只有64GB，没有办法将所有的图片都预加载到内存中，所以由于内存容量不够导致我们不得不放弃数据预加载到内存中的方式。

1.2 数据预先处理，存储在磁盘中-受限于磁盘IO速度

所以由于数据集数目巨大，在Pytorch中写Dataloader时，我先在__init__函数中获取所有的图片的路径，然后在__getitem__函数中根据图片路径读取图片数据，并进行数据增强，这种在模型训练的过程中读取文件并从内存发送到GPU进行训练的方式极大地依赖磁盘的读写速度，如果磁盘的读写速度很差，那么GPU将受限于IO瓶颈，并不能发挥GPU的全部性能。

我先将这100W+的图片数据集放在了一个机械硬盘中进行训练，硬盘是希捷酷鱼4T 7200转，设置Dataloader的参数为num_workers=8，pin_memory=True，batch_size=256，GPU为3090 24G，在训练的过程中，GPU的利用率在长时间段为0%，短时间跳动到7%，这就是典型的GPU在等待CPU从磁盘中读取数据的表现。

为了进行比较，我之后将这100W+的图片数据集放在固态硬盘三星 980 1T上，Dataloader的参数与上述保持一致，GPU同样为3090 24G，GPU的利用率长期保持在20%，训练速度得到提升。

1.3 小结

从上述文章中，我们知道，影响模型训练速度不仅仅是跟你的GPU硬件相关，如果在你的GPU的速度足够好的情况下，如果训练速度还是达不到，那么就需要从内存频率，磁盘IO速度，CPU负载上进行排查，上述的内存频率、磁盘IO速度、CPU的处理速度都可能是限制你训练速度上不去的原因，所以说深度学习需要各个硬件在全方位上进行匹配，以达到更快的训练速度。