[3D数据深度学习] （PC/服务器集群cluster）CPU内存/GPU显存限制及解决办法

3D数据的深度学习目前研究远不如2D深度学习成熟，其中最大的一个原因之一就是收到硬件条件的限制。3D数据虽说只比2D数据增加了一个维度，但所占据的内存却是成倍的增长。对于3D数据的深度学习，我们会分析其在CPU内存和GPU显存两方面的限制，希望大家能够充分利用自己的资源进行深度学习。

1. 硬件配置推荐

CPU: 大内存，多核（很关键，越多越好）高性能CPU
GPU: 大显存（24G以上），比如A6000（48G），TeslaV100（32G）
服务器集群： 每个节点多CPU，大内存，多卡（GPU）

2. 深度学习流程及遇到的问题

在深度学习代码中，以pytorch为例，CPU往往负责着数据读取/加载，数据增强的部分。而GPU负责模型模型的训练。
在CPU上完成数据读取，数据增强之后，将数据传到GPU上进行训练。所以我们的数据要同时满足CPU和GPU的限制。在执行程序时，我们既不能超过CPU内存的限制（否则会发生线程中断，程序卡死），也不能超过GPU显存的限制（否则会发生OOM）。但同时，我们也希望尽可能最大化的利用内存和显存，提高我们的训练速度，特别对于3D数据，一般只能使用较小的batch_size, 所以训练时间往往非常长，正确的参数设置可以大大提高我们的训练时间。下面我会像大家介绍如何在不超过自己内存/显存限制的情况下优化参数配置，加速训练。

3. CPU内存限制及参数设置

首先我们要知道，在进行训练时，我们的cpu并不是每次一读取batch个数据进行网络更新，再读取下一个batch数据进行训练。这里有一个关键参数：num_workers Pytorch dataloader中的num_workers (选择最合适的num_workers值)

dataloader一次性创建num_worker个worker，（也可以说dataloader一次性创建num_worker个工作进程，worker也是普通的工作进程），并用batch_sampler将指定batch分配给指定worker，worker将它负责的batch加载进RAM。

num_worker设置得大，好处是寻batch速度快，因为下一轮迭代的batch很可能在上一轮/上上一轮…迭代时已经加载好了。坏处是内存开销大，也加重了CPU负担（worker加载数据到RAM的进程是CPU复制的嘛）。num_workers的经验设置值是自己电脑/服务器的CPU核心数，如果CPU很强、RAM也很充足，就可以设置得更大些。

num_worker小了的情况，主进程采集完最后一个worker的batch。此时需要回去采集第一个worker产生的第二个batch。如果该worker此时没有采集完，主线程会卡在这里等。（这种情况出现在，num_works数量少或者batchsize 比较小，显卡很快就计算完了，CPU对GPU供不应求。

num_works数量越多，CPU准备数据越快 pytorch dataloader 使用batch和 num_works参数的原理是什么？
红线情况：数据很大，num_works很少GPU计算完成后想要拿到下一个批次的数据时发现CPU还没准备好数据。要等待很久才能开始下一次计算，从而影响到GPU使用率。是3D数据训练时常见到的情况。

最好的情况是紫线这种：GPU计算的过程中CPU已经准备好数据了。等GPU计算完上一个批次数据，直接把下一个批次数据发给GPU。

---

由于3D数据很大，我们的batch_size因为显存的限制往往只能设置的很小（比如2，4，6）所以网络在收到小batch_size个数据之后往往很快就完成了训练，等待下一批batch数据到来。因此，对于体积较大的3D数据，其时间限制主要来自于数据读取部分，因此，数据的并行读取比较重要，我们希望使用尽量多个CPU同时读取数据。也就是说，在不超过我们的内存限制的情况下，num_worker越大越好。

对于PC，首先我们需要查看自己的硬件配置，有多少个CPU，num_worker不能超过CPU个数。
对于服务器，多个脚本程序往往在一个节点上执行，这些程序共享这个节点的内存和CPU，所以我们要合理安排每个程序的CPU使用数量以及内存占用率。如果是特别大的3D数据，建议一个节点只执行这一个程序，来充分利用CPU资源

我们可以使用htop命令来查看自己的CPU个数以及内存使用情况

如图，我一共有64个CPU，内存最大为188G
理论上说，我的num_worker最大可以设置为64，但实际上，当64个CPU同时加载数据，很容易超出188G内存的限制，如果我们使用服务器的话，可以使用以下命令对每个CPU所使用的最大内存做出限制：

#SBATCH --mem-per-cpu=3G

180G内存/60个CPU = 3G，当每个CPU只使用3G内存，就不会超过最大内存了。

然而，当我们的数据非常大，我们就需要降低num_worker的个数来避免超出内存，比如说我们的原始数据大小为（834，224，834，3），在多个CPU读取数据时会占用特别大的内存，所以我们要测试num_worker对于不同数据的最大值

经过测试，对于（834，224，834，3）的3D数据，num_workers最大值为18，如果再大的话多个CPU读取数据产生的峰值会超过内存限制，从而引起程序中断

4. GPU显存限制及参数设置

对于3D数据训练，超过GPU显存最大值（OOM）是最常遇到的问题。减小Batch_size是简单的做法，但是很多情况下Batch_size已经很小了，没法再小。

下面是三种解决办法，根据不同的情况选择：

1. DataParallel: 模型不算很大，batch_size 较小时能够运行没有OOM，但效果不好，希望增大batch_size进行训练。可以使用模型并行扩大batch_size, 使用4张卡同时训练的话可以将batch_size扩大4倍。
2. Model parallel:模型非常大，一张卡放不下，可以将模型的不同部分放到不同卡上进行训练。适用于多支路网络，将不同支路放到不同卡上，不会对训练时间造成太大影响。若是单支路网络使用模型并行，训练时间会变长。
3. Gradient checkpointing：模型非常大，一张卡放不下，但是只有一张卡。将模型分段保存，在训练时会分段进行训练，会明显增长训练时间，预算有限的情况下可以使用。

服务器集群上的DataParallel详情见我之前的教程：[pytorch] 分布式训练 Distributed Data-Parallel Training (DDP)
模型并行Model parallel的教程会在以后发

因为batch_size很小，所以各种加速训练的技巧表现都不是很好，我们只讨论OOM的解决办法