基于文件存储UFS的Pytorch训练IO五倍提升实践

我们在协助某AI客户排查一个UFS文件存储的性能case时发现，其使用的Pytorch训练IO性能和硬件的IO能力有很大的差距（后面内容有具体性能对比数据）。

让我们感到困惑的是：UFS文件存储，我们使用fio自测可以达到单实例最低10Gbps带宽、IOPS也可达到2w以上。该AI客户在高IOPS要求的AI单机小模型训练场景下，或者之前使用MXNet、TensorFlow框架时，IO都能跑到UFS理论性能，甚至在大型分布式训练场景中，UFS也可以完全胜任。

于是我们开启了和客户的一次深度联合排查。

初步尝试优化

1、调整参数

基于上述情况，首先考虑是不是使用Pytorch的姿势不对？参考网上提到经验，客户调整batch_size、Dataloader等参数。

Batch_size

默认batch_size为256，根据内存和显存配置尝试更改batch_size大小，让一次读取数据更多，发现实际对效率没有提升。通过分析是由于batch_size设置与数据读取逻辑没有直接关系，IO始终会保留单队列与后端交互，不会降低网络交互上的整体延时（因为用的是UFS文件存储，后面会讲到为什么用）。

Pytorch Dataloader

Pytorch框架dataloader的worker负责数据的读取和加载、分配。通过batch_sampler将batch数据分配给对应的worker，由worker从磁盘读取数据并加载数据到内存，dataloader从内存中读取相应batch做迭代训练。这里尝试调整了worker_num参数为CPU核数或倍数，发现提升有限，反而内存和CPU的开销提升了不少，整体加重了训练设备的负担，通过 worker加载数据时的网络开销并不会降低，与本地SSD盘差距依然存在。

这个也不难理解，后面用strace排查的时候，看到CPU更多的时候在等待。

所以：从目前信息来看，调整Pytorch框架参数对性能几乎没有影响。

2、尝试不同存储产品

在客户调整参数的同时，我们也使用了三种存储做验证，来看这里是否存在性能差异、差异到底有多大。在三种存储产品上放上同样的数据集：

1、 单张平均大小20KB的小图片，总量2w张。
2、 以目录树方式存到三种存储下的相同路径，使用Pytorch常用的标准读图接口CV2和PIL

测试结果，如下图：

读取方式	SSHFS	本地SSD	UFS文件存储
CV2	319.94张/s	554.73张/s	72.41张/s
PIL（image_open）	435.93张/s	3507.93张/s	115.78张/s

注：SSHFS基于X86物理机（32核/64G/480G SSD*6 raid10）搭建，网络25Gbps
结论：通过对存储性能实测， UFS文件存储较本地盘、单机SSHFS性能差距较大。

为什么会选用这两种存储（SSHFS和本地SSD）做UFS性能对比？