PyTorch 训练时中遇到的卡住停住等问题

前言

在实际训练中遇到了各种各样的卡住问题，在此总结一下，PyTorch 训练时遇到的卡住停住等问题可以从以下几个方面根据情况具体分析 (参考PyTorch训练时，Dataloader卡死、挂起，跑一个epoch停了，问题解决方案)：

• 前一进程还未处理完，又进入下一个导致互锁：在每个Epoch后，或每个 batch 后暂停一下： time.sleep(0.003)；
• 内存问题：使用开关，pin_memory=True/False
• 多进程导致互锁问题：减少进程数或者不用，num_workers=0/1
• OpenCV与Pytorch互锁；
• batch size 的设置；

1. PyTorch 训练时在第一个 epoch 的最后一个 batch 卡死

- 问题描述

使用 PyTorch 框架训练模型，训练第一个 epoch 时，在最后一个 batch 处卡死，卡了一天都没有动弹，而 CPU 和 GPU 都处于正常运行的状态，程序也没有报错，并且之前训练一直都是正常的。最终，只能通过 Ctrl+C 强制性暂停。如下图所示。

- 可能的原因

1. 搜索文章发现，有人出现这种问题是和 cv2.imread 有关，用 OpenCV 的接口进行数据读取，而没有用 PIL，导致出现 OpenCV与Pytorch互锁的问题，关闭OpenCV的多线程即可解决问题^{1 2}。
2. 但是我们的代码中并没有使用 OpenCV，通过 Debug 发现，出现问题的时候，最后一个 batch =1，而我们使用的是四块 GPU 进行训练，原因就在此。

- 解决方法

Pytorch 的多 GPU 处理接口是 torch.nn.DataParallel(module, device_ids)，具体并行机制可参考：Pytorch多GPU并行处理。
该接口还要求输入数据的 batch 数量要不小于所指定的 GPU 数量。另根据官网的解释和注释 (The batch size should be larger than the number of GPUs used.)，batch的数量会均分到每块GPU上进行处理，因此要保证一个整数的关系。

故，一定要注意在使用多块 GPU 训练时，注意 batch_size 的取值，避免出现最后一个 batch 的实际size小于所指定的 GPU 数量的情况。

多 GPU 并行训练中其他比较坑的地方可参考：pytorch多GPU数据并行模式 踩坑指南，和nn.DataParallel说再见。

2. 使用命令行运行时，卡在第一个 epoch

- 问题描述

使用 PyTorch 框架训练模型，在 Pycharm 里面可以正常运行，但是在命令行中运行时，在第一个 epoch 处卡死。

- 原因分析

经过定位，发现卡死的地方是PyTorch自带的 DataLoader 处有问题， 使用batch和 num_works参数的原理可参考此文³，DataLoader的函数定义如下⁴：

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
num_workers=0, collate_fn=default_collate, pin_memory=False,
drop_last=False)

其中 num_workers 表示使用多进程加载的进程数，0代表不使用多进程，pin_memory 表示是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些。

- 解决方法

把 num_workers 改为 0 即可正常运行，如何设置还是要靠具体实验。

---

1. Pytorch dataloader在加载最后一个batch时卡死 ↩︎

2. 在PyTorch训练一个epoch时，模型不能接着训练，Dataloader卡死 ↩︎

3. pytorch dataloader 使用batch和 num_works参数的原理是什么？ - 孟paper的回答 - 知乎 ↩︎

4. pytorch DataLoader num_workers 出现的问题 ↩︎