记一次深度学习数据传输及训练的性能优化

开端

某校生物信息医学工程博士生（下称博士）重金邀请，让我辅助他进行针对 LUNA 肺癌数据集的结节 3D 分割模型的构建。说是重金，其实给的报酬也不是很多，然而盛情难却，加上和这个博士关系不错，还是接了这一手生意。

在邀请我进行辅助之前，博士已在学校的计算节点上耗费将近一星期时间，将所需的训练/测试集的三维图像、标签从 .mhd/.raw 文件处理为 .npy 格式的 numpy 矩阵文件。数据集有了，模型有现成的作参考，博士使用 PyTorch 框架，不到半个小时就无比熟练地搭建出了一个 3D UNet 的模型。简单调试，证明网络没有明显问题后，下面就准备开始进行正式的训练了。

海量数据集的困惑

由于用于处理数据的服务器与实际训练的服务器不在同一台上，我们需要对处理过的数据集先进行转移。简单的 du -hs 命令发现，区区 2000 张的训练集、验证集的三维灰度图像、二进制标记文件的总大小在 .npy 格式下就已经达到了惊人的 1TB 的大小。博士吃惊之余，还是随手打开了 WinSCP，开始把文件从训练服务器向测试服务器拷贝。

博士点燃一根烟起身离开，我抿了一口保温杯里的水，盯着屏幕上 100MB/s 的“神速”不禁陷入了沉思。照这个速度下去，1024 GB × 1024 GB/MB ÷ 100 MB/s ÷ 3600s/h，要拷完这些数据集得花上三个小时左右。有没有什么办法能加速这个过程呢？

压缩与围观

我的第一想法就是压缩。因为在实际进行深度学习的过程中，.npy 格式的数据集实属罕见，能见到的比较接近的范例是 .npz 格式，也就是经过压缩的 .npy 格式。并且在这个肺部的三维的灰度立体图像中，显然数据集本身属于稀疏矩阵，而此时压缩的效果必然非常明显。

说干就干，我不等博士回来，直接把 WinSCP 的拷贝过程停掉，然后登录到了数据预处理的服务器上。这是一台对称双 Intel Xeon 处理器的 32 核心 64 线程的服务器，配有 256 GB 的服务器内存，可谓性能拔群。根据 WinSCP 上指示的路径，我定位到处理后的数据集所在的目录，执行 tar -cvzf ../dataset.tar.gz . 命令，开始对数据集进行压缩。

屏幕上不断刷出一行一行的数据集文件的完整路径，盯着屏幕我又陷入了思考。调查了一下数据集文件的总数目，然后根据屏幕滚屏的速度，大致估算出了压缩这些数据集所需要的时间——2个小时。这显然是不可接受的，我的心变得有些纠结了起来：是不是应该重新把 WinSCP 的拷贝过程再点开一下？还是再挣扎一下看看？

处于本能，我在服务器上习惯性地输入了 htop 命令查看系统的负载情况，没想到这一看就看出了问题。经典界面映入眼帘，服务器只有一个核心在 100% 运作，剩下的所有核心的负载程度都不足 1%，孤零零的满载核心非常扎眼。往下扫视，CPU 占用第一的程序便是我们的压缩程序 tar。这可当真应了那句主机界的话：一核有难，六十三核围观。

并行压缩与奇迹

到了这个地步，我们不难意识到这样一个问题：传统的 tar 压缩程序是单线程执行的，而在我们的服务器上，根本没有充分地利用服务器庞大的算力资源，导致了“线程围观惨案”，也导致了完成时间不可接受。

博士抽完烟回来了，看着我神奇的操作，刚想问我发生了什么就被我示意堵住了嘴巴。我在博士的电脑上打开了 Google，然后搜索 linux parallel compression 这样的关键字。然而刚想按下回车键，我突然想起在 Arch Linux 上安装 docker 的时候的几个可选依赖，其中有一个可选依赖叫做 pigz，其中的说明中有相关的字样。

说干就干，在博士的注目下我掏出了随身的笔记本，执行 pacman -Si pigz：

软件库         : community
名字           : pigz
版本           : 2.4-1
描述           : Parallel implementation of the gzip file compressor
架构           : x86_64
URL            : https://www.zlib.net/pigz
......

其中的描述赫然就写着“并行实现的 gzip 文件压缩器”字样。使用 man pigz 命令仔细查阅一番用户手册之后，我使用 tar -cvf - . | pigz -9 -p 60 > ../dataset.tar.gz 命令，同时开启 60 个线程进行并行压缩。

回车键甫一按下，屏幕立即就像要爆炸了一般疯狂刷出日志。我停掉进程，然后把他挂在一个 screen 会话中继续运行。不到十分钟的功夫，当我上去再次检查的时候，发现压缩操作已经完成了，最终得到了一个只有 3.9GB 的单个文件。

这是几乎令人疯狂的压缩率，接近 300:1 的压缩比率，使得整个数据集从预处理服务器移动到训练服务器最终只花费了 40 秒左右的时间，而解压操作只花费了不到五分钟的时间。操作完成之后，博士整个人都看呆了。他本人表示，原本打算打两盘 Dota 2 打发一下时间来着，没想到这么快就把数据集传完了。

训练与交织的灯光

用于训练的服务器（其实只是放在博士实验室里的一台高性能主机）上装有 Intel i7-8700K，据称改装了液金和水冷之后已经稳定超频到了 5.3 GHz；用于图形与计算加速的 AORUS 1080Ti 据称也解锁了 TDP 限制、改掉了频率，功耗直奔 500W。不过这些的一切都不重要，重要的是博士在我刷手机的时候已经把代码用U盘拷贝到训练主机上，并且成功地把代码跑起来了。

震耳欲聋的冷排风扇加速声把我的视线拉向了实验主机。机箱前盖板的硬盘灯以人眼几乎无法捕捉到在闪动的高速在“连续”地发出刺眼的红色光芒，映衬着白色灯光的美商海盗船冷头显得格外迷人。几分钟之后，训练程序正式开始调用显卡进行计算加速了，我甚至能看到显卡上的风扇指示灯在有规律地一跳一跳地闪烁，连带着机箱的开关电源随之发出不太响亮的啸叫声，交织着满耳的希捷硬盘声，整个机箱仿佛是在演奏一曲深度学习的赛博朋克。

一个季节的训练

无聊至极，按下 Ctrl+Alt+F2 切换到另一个 tty，仔细看了一下他的代码。训练代码本身写的非常不错，还使用了 tensorboard 作为可视化工具，但是没有注意到他是否开启了 tensorboard；除此之外，他还在训练的循环中使用了 tqdm 以追踪训练的进度、预估训练所需的总时间。

代码本身并不长，草草看完之后便失去了兴趣。在询问了是否需要开启 tensorboard 并得到了肯定的回答、随手把 tensorboard 挂上之后，切换回训练代码所在的 tty，tqdm 的输出却让我大吃了一惊：屏幕上的进度条告诉我，这个训练程序向神经网络里输入 1 张图像需要 3.8 秒，而总共需要训练 1000 个 epoch，如果这样下去，包括验证集的验证过程，总时长将达到……

谁是瓶颈？

嗯，没错。依这个速度算，估计要从今年跑到明年开春之后才能跑完。

我把这个事情立即告诉了博士，博士本人表示也非常无奈：本来三维图像就非常吃显存，原图都是 512 x 512 x 512 这种尺寸的，就算每个像素点都用 float32 去存储，最好情况下理论上每张图也要吃掉半个G的显存，实际训练的时候可能连一张图都装不进去。博士表示，为此，他在代码中使用 skimage 的 resize 方法，在训练的 DatasetLoader 里使用立方插值将数据 resize 成 256 x 256 x 256 的尺寸，再放进神经网络里训练，这才堪堪让显存够用。

听到这里，我突然想到了一个非常可怕的问题：三维图像的 resize 操作是非常吃 CPU 资源的，尤其是这么大型的三维图像的操作，同时还要进行立方插值操作。很有可能这个程序本身的瓶颈并不在显卡上，而是在从硬盘读取超大的数据集载入内存，然后在内存中进行超大图像的 resize + interpolation 操作这两个步骤上。身旁主机间隔几秒钟狂灭一下的 FAN STOP 指示灯和电源的啸叫声也提醒了我，这段代码绝对是有问题的。

想明白之后，我又掏出了 htop 工具查看系统负载情况。实际上，系统的负载并不高，但是 iowait 值却的确比较高，这说明硬盘可能是潜在的性能瓶颈点，光彩熠熠的红色硬盘灯与扎耳的硬盘声音似乎也在向我抱怨。

但我还是不能太早得出结论。我查看了数据集的 .npy 文件，发现大部分文件的大小能达到 1～2 GB，硬盘的读取速度可能是瓶颈的主要成因，但应该不是主要问题。检查代码的 DataLoader 发现，博士为了充分利用主机性能，将 DataLoader 的 num_workers 参数设为 4，从而可能导致了硬盘时刻处于满载状态。目前为止，似乎所有的瓶颈都指向了硬盘。

但我并不关心到底谁是瓶颈，因为他们都不重要。在我看来，目前为止，问题的本质在于预处理过程做的不够充分。试想如果不需要从硬盘载入一个上 GB 大小的三维图像，然后再执行极为消耗计算资源的 resize + interpolation 操作，这个性能瓶颈本不应该出现。如果提前在预处理中将图像缩放为适合的大小，不仅传输的文件尺寸会减小很多，计算速度也会有质的飞跃。

二次预处理：不再围观

我将我的想法和出现当前现象的成因假设解释给博士听，博士恍然大悟。在我的注视下，博士将训练代码的预处理代码提取出来，写成一个新的数据预处理 Python 程序，然后打算放到服务器上去执行。博士的代码没有什么大问题，但是刚刚进行的一系列优化操作让我突然警醒：如果直接单线程去运行这样的程序，是否还会出现“一核有难、六十三核围观”的惨状呢？

我将我的想法与博士分享之后，博士将电脑键盘转交给我然后坐在了一边。我使用 threadpool 简单地封装了一下博士的处理过程，然后将并发线程数设置为 60，将代码上传到预处理服务器上，启动一个新的 screen 会话，然后运行我们的二次预处理程序。

程序运行的很快，并且这次在 htop 中看的时候，所有核心基本都处于满载的状态，界面中被彩色条纹填充得满满的框框证明了我们的代码充分地利用了这个服务器的计算资源。果不其然，不消二十分钟，我们就跑完了第二次预处理。

训练出毛病了？

使用同样的 pigz 并行压缩、拷贝文件、解压文件，我们把二次预处理的数据集转移到了训练服务器上。更新了训练代码之后，我们再次启动训练程序。随着回车键的按下，显卡的风扇转速指示灯应声而灭，随之而来的是显卡的风扇就像飞机发动机一样高速飞转了起来。硬盘灯间歇地一闪一闪，已没有刚开始训练时那样密集；CPU 的冷排风扇也开始变得温柔起来。一切是那么地和谐。

回观屏幕，我们的训练速度已经达到了每秒 3 个 batch 的飞速，而每个 batch 的 size 是 4，比起之前的代码提升了大约 30 倍的速度。从一个季节的漫长等待优化到接近半个星期就能出结果，很难想象我们刚刚完成了这样的壮举。

但是，事实是无比残酷的。正当我们飞速跑完了训练过程的第一个 epoch，完成验证集的演算并显示 metrics 在屏幕上的时候，我和我的小伙伴们都惊呆了：这次训练的 training accuray 和 validation accuracy 竟然高达 0.00！

博士当时就慌了起来，我仔细思考后，建议博士再观察一个 epoch 再下定论。由于我们数据集的目标结节区域的体积占整个三维空间的大小是非常小的，对于一个 256x256x256、体积为 16777216 像素块的三维空间，很有可能只有其中的 16x16x16、体积为 4096 像素块的区域是目标的标定区域。那么，单纯地使用传统的 accuracy 来评估模型的效果很可能是不太实际的，而这里第一轮训练完成后出现的 0.00 accuracy 情况可能也是合理的。

向博士解释完可能的原因之后，博士一拍脑门想起来了 dice loss 可能比较适合于这种不平衡的数据集的量度，于是溜回自己的电脑前开始面向 GitHub 编程，在训练代码中实现 dice loss，我则坐在风扇轰鸣、流光溢彩的训练机前盯着屏幕不断向前滚动的第二轮训练的进度条发呆。

很快，第二轮训练完成了。在等待训练进行的百无聊赖之余，我切换到另一个 tty，调出 python，加载了一个标签数据，打算进行进一步的探索。直接 print 标签所得到的输出是经过 numpy 自动修剪之后得到的输出，毫无疑问，标签数据的四边全是 0。思考之后，我使用了 np.sum 对标签数据直接进行求和，令人震惊的事情出现了：这个标签求和的结果是 0！

看着正在编码的博士，我决定仔细探索之后再向他说明问题。因为我还不能确定是处理过之后的数据集出现了问题，对于阴性样本来说，数据集的标签全为 0 才是正常的。简单使用 os.listdir 遍历了一下标签目录，然后对所有的标签分别进行求和操作，结果非常令人失望——所有的标签数据都是 0！

我立刻切换 tty，按住左下角的 Ctrl 键就狂砸 C 字键帽，雷柏国产青轴被我打得咯咯作响。博士看起来刚刚写完 dice loss 的代码，正要和我交流整合 dice loss 的事情，就看到我把训练直接停掉了，忙询问我原因。知晓情况后，博士沉默了，表示自己需要冷静一下，之后掏出打火机和烟，再度出门思考人生。

这显然是博士的预处理程序有问题，将正常的标签处理成了全为 0 的无效标签。单步调试发现，博士在使用 skimage 的 resize 方法时，原标签的 0-1 表示由于具有 np.uint8 类型，可能被当成了 0-255 的灰度图像，进而在进行插值、缩放时导致了所得的目标值小于 1，从而导致了这部分值被舍去变成 0。因为之后在显卡上的训练需要将数据集再转成 np.float32 类型，博士在原代码的最后才进行这个转型操作，而这时候再进行转型已经没有任何意义。

将转型操作提前到 resize 之前重试单步运行，发现 resize 还是不能自动确认原图像像素值的上下界，问题依旧。我将保温杯里的水一饮而尽，思考过后想出了一个 workaround：先将像素值转型为 np.float32，将像素值整体乘以 255 后丢进 resize，得到 0～255 的像素值集；之后对像素值进行二值化处理，将大于 127 的变为 1.0，将小于等于 127 的变为 0.0 即可。

博士吸烟归来，验证 workaround 及 dice loss 的有效性之后，我们将代码整合进原有的训练中，再度开启训练。一个 epoch 跑完之后，dice loss 和 accuracy 的值终于回归了正常，博士的脸上也终于有了微笑。

后续

模型训练成功，但是由于博士先前超参数设置过于随意、没有设置 checkpoint 等原因，最后得到的模型过拟合非常严重。在训练集上的准确率达到了 99.6%，但在验证集上的准确率不足 80%。就此一事，博士表示自己对基础设施的运用还是 too young & sometimes naive，今后必须要学习一个，涨涨自己的姿势水平。我也在此期间收获不少。

最后，博士并没有兑现他给我开出的报酬，只是请我吃了顿饭。但知识和精神上的满足，远比吃进肚子里的过客要更令人愉悦。