Pytorch 训练时卡在dataloader

Pytorch 训练时卡在dataloader

1. 问题描述

Traceback (most recent call last):
  File "train_cross_dataset.py", line 281, in <module>
    srocc_all[i], plcc_all[i] = solver.train()
  File "train_cross_dataset.py", line 129, in train
    img, label = next(iter_list[idx])
  File "/home/server01/anaconda3/envs/pytorch/lib/python3.8/site-packages/prefetch_generator/__init__.py", line 92, in __next__
    return self.next()
  File "/home/server01/anaconda3/envs/pytorch/lib/python3.8/site-packages/prefetch_generator/__init__.py", line 85, in next
    next_item = self.queue.get()
  File "/home/server01/anaconda3/envs/pytorch/lib/python3.8/queue.py", line 170, in get
    self.not_empty.wait()
  File "/home/server01/anaconda3/envs/pytorch/lib/python3.8/threading.py", line 302, in wait
    waiter.acquire()
KeyboardInterrupt

可以看出训练过程卡在从dataloader中提取数据的过程。

2. 解决方法

（1）Dataloader里面不用cv2.imread进行读取图片，用cv2.imread还会带来一系列的不方便，比如BGR通道问题，以及不能使用torchvision进行数据增强之类的操作，一般使用PIL包里面的image.open()来读取图片。（不适用于本例）

（2）将Dataloader里面的参变量num_workers设置为0，但会导致数据读取很慢，拖慢整个模型的训练。

（3）如果用了cv2.imread，不想改代码的，那就加两条语句，来关闭Opencv的多线程：cv2.setNumThreads(0)和cv2.ocl.setUseOpenCL(False)。加了这两条语句之后，并不影响模型训练。（不适用于本例）

（4）这种情况应该是属于pytorch多线程锁死，在github上看到有该问题，但是没有解决的。

3. 参考建议

首先确保num_works数量低于CPU数量(如果使用Kubernetes，则设置为pod)，但是设置得足够高，使数据随时可以用于下一次迭代。如果GPU在t秒内运行每个迭代，而每个dataloader worker加载/处理单个批处理需要N*t秒，那么您应该将num_workers设置为至少N，以避免GPU停滞。当然，系统中至少要有N个cpu。

不幸的是，如果Dataloader使用任何使用K个线程的库，那么生成的进程数量就会变成num_workersK = NK。这可能比计算机中的cpu数量大得多。这会使pod节流，而Dataloader会变得非常慢。这可能导致Dataloader不返回批处理每t秒，导致GPU暂停。

避免K个线程的一种方法是通过OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 python train.py调用主脚本。这就限制了每个Dataloader工作程序只能使用一个线程，从而避免了使机器不堪重负。你仍然需要有足够的num_workers来满足GPU的需要。

您还应该在__getitem__中优化您的代码，以便每个worker在较短的时间内完成其批处理。请确保worker完成批处理的时间不受从磁盘读取训练数据的时间(特别是当您从网络存储中读取数据时)或网络带宽(当您从网络磁盘读取数据时)的影响。如果您的数据集很小，并且您有足够的RAM，那么可以考虑将数据集移动到RAM(或/tmpfs)中，并从那里读取数据以进行快速访问。对于Kubernetes，您可以创建一个RAM磁盘(在Kubernetes中搜索emptyDir)。

如果你已经优化了你的__getitem__代码，并确保磁盘访问/网络访问不是罪魁祸首，但仍然会出现问题，你将需要请求更多的cpu(为了一个Kubernetes pod)，或者将你的GPU移动到拥有更多cpu的机器上。

另一个选项是减少batch_size，这样每个worker要做的工作就会减少，并且可以更快地完成预处理。后一种选择在某些情况下是不可取的，因为会有空闲的GPU内存不被利用。

你也可以考虑离线做一些预处理，减轻每个worker的负担。例如，如果每个worker正在读取一个wav文件并计算音频文件的谱图，那么可以考虑离线预先计算谱图，只从工作者的磁盘中读取计算的谱图。这将减少每个worker的工作量。
你也可以考虑将dataloader里的设置pin_memory=False。

上述的方法来自here
参考链接：https://blog.csdn.net/qq_32526087/article/details/106350530