pytorch dataloader数据加载占用了大部分时间，各位大佬都是怎么解决的？

这问题太真实了！DataLoader 慢，简直是训练时候的拦路虎。看着 GPU 闲着干等数据，心里那个滋味，简直酸爽。网上关于 DataLoader 优化的文章不少，但很多都泛泛而谈，或者没说到点子上。作为摸爬滚打了好一阵子的人，想跟各位分享下我自己实践下来比较靠谱的几招，希望能帮大家少走弯路。

核心思路：把数据加载这件事，从 CPU 往 GPU 身边“挪”

咱们的 GPU 是个吃“生食”的大家伙，它最快，能处理大量数据。而 CPU 呢，更擅长精细活，比如文件读取、预处理、数据增强什么的。DataLoader 的瓶颈，往往出在 CPU 忙不过来，没法及时把处理好的数据喂给 GPU。所以，我们的目标就是让 CPU 尽量提前把数据准备好，而且准备好了就赶紧传给 GPU。

第一招：给 DataLoader 加足马力——`num_workers` 的秘密

这是最基本，也是最直接的优化方法。

怎么回事？ `num_workers` 参数决定了有多少个子进程（worker processes）在并行加载和预处理数据。默认是 0，意味着主进程自己干所有活，那肯定慢。
怎么调？
经验值： 一般建议从 CPU 核心数 / 2 开始尝试，然后逐渐往上加。比如你 CPU 有 8 个核心，可以先试试 `num_workers=4`。
不要太贪心： 并不是越多越好。如果 `num_workers` 设得比 CPU 核心数还多，子进程之间会为了争夺 CPU 资源而互相干扰，反而会降低效率。甚至可能导致死锁或者内存占用爆炸。
根据数据集和预处理复杂度： 如果你的数据预处理非常耗 CPU（比如复杂的图像增强、文本 tokenization），你可能需要更多的 workers。如果预处理很简单，或者主要瓶颈在于磁盘 I/O，workers 的数量增加可能效果不那么明显。
内存考量： 每个 worker 都会占用一部分内存，加载和预处理数据。如果你的数据集很大，或者预处理过程中会产生很多中间变量，需要确保你的机器有足够的内存来支撑设定的 `num_workers`。我见过不少因为 workers 太多导致 OOM（Out Of Memory）的，那可就得不偿失了。
如何检查效果？
观察 GPU 利用率： 训练时，用 `nvidiasmi` 或者 TensorBoard 的 GPU 监控面板，看看 GPU 的利用率是不是能跑到 90% 以上。如果老是跳水，那说明 DataLoader 没跟上。
计时： 简单粗暴点，自己用 `time.time()` 计时看看加载一个 batch 的平均时间。

第二招：数据预处理前置——`pin_memory=True`

这个参数是 PyTorchDataLoader 自带的“小聪明”，但效果奇佳。

怎么回事？ 当 `pin_memory=True` 时，DataLoader 会在加载数据时，把数据所在的内存“固定”住（pinned memory）。这样做的好处是，当数据被传输到 GPU 显存时，这个过程会更快，因为它避免了从普通内存（host memory）到显存（device memory）的额外拷贝。
怎么用？ 直接在 `DataLoader` 初始化时加上 `pin_memory=True` 即可。
```python
train_loader = DataLoader(
train_dataset,
batch_size=32,
shuffle=True,
num_workers=4,
pin_memory=True < 加上这句！
)
```
注意事项：
与 `num_workers` 配合： `pin_memory=True` 只在 `num_workers > 0` 时才生效。如果 `num_workers=0`，它就没有作用。
对显卡有要求： 理论上，这个特性对 PyTorch 和 CUDA 的版本有一定要求，但现在主流版本都没问题。
内存占用： 同样，pinned memory 也会占用一部分内存，但通常比创建大量子进程的内存开销要稳定。

第三招：充分利用 GPU 的“预加载”能力——`prefetch_factor`

这个参数相对 `num_workers` 来说，是更细粒度的控制。

怎么回事？ `prefetch_factor` 控制了每个 worker 进程会提前加载多少个 batch 的数据。它和 `num_workers` 配合工作。简单来说，DataLoader 会维护一个队列，里面是 worker 进程已经加载好并准备好传输给 GPU 的数据。`prefetch_factor` 就是这个队列的“预读”深度。
怎么调？
经验值： 默认值通常是 2。一个比较常用的值是 2 到 5。
与 `num_workers` 的关系： `prefetch_factor` 的作用是，当 worker 进程把一个 batch 数据准备好后，它不会立刻停止，而是会接着准备下一个 batch，直到达到 `prefetch_factor` 的数量。这样可以更充分地利用 worker 进程和 CPU 的计算能力。
权衡： 设得太高，可能会占用更多内存。设得太低，又可能无法完全填满 GPU 的“饥饿感”。
如何结合使用？
```python
train_loader = DataLoader(
train_dataset,
batch_size=32,
shuffle=True,
num_workers=4,
pin_memory=True,
prefetch_factor=2 < 尝试调整这个值
)
```
建议： 先尝试 `num_workers` 和 `pin_memory=True`，如果 GPU 利用率还是不够高，再考虑调整 `prefetch_factor`。

第四招：数据缓存与预加载——自己的 DataLoader 包装

有时候，即使上面几招都用了，DataLoader 依然慢。这时，可能需要更深入地处理数据加载流程。

怎么回事？
1. 内存缓存： 如果你的数据集虽然大，但可以全部或者大部分放入内存，可以考虑在程序启动时，一次性把所有数据（或一大部分）加载到 Python 的 list 或 NumPy 数组中。然后，你的 `__getitem__` 方法就只需要从内存中提取数据，这比频繁地从磁盘读取要快得多。
```python
class MyCachedDataset(Dataset):
def __init__(self, data_path):
self.data = self._load_all_data(data_path) 这是一个耗时的操作，但只执行一次

def _load_all_data(self, data_path):
这里实现你的数据加载逻辑，比如读取所有图片文件路径，
然后预处理并缓存到内存中 (e.g., list of numpy arrays, torch tensors)
print("Loading and caching all data into memory...")
cached_data = []
... your data loading and preprocessing ...
return cached_data

def __len__(self):
return len(self.data)

def __getitem__(self, idx):
return self.data[idx]

使用时
dataset = MyCachedDataset("path/to/your/data")
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)
```
2. 自定义预加载线程： 对于非常大的数据集，或者预处理非常复杂，无法完全放入内存。你可以自己实现一个后台线程，专门负责加载和预处理数据，并将处理好的 batch 存入一个队列（比如 `queue.Queue`）。主训练循环则从这个队列中获取 batch。
```python
import threading
import queue

class BackgroundDataLoader:
def __init__(self, dataset, batch_size, shuffle, num_workers, prefetch_size=10):
self.dataset = dataset
self.batch_size = batch_size
self.shuffle = shuffle
self.num_workers = num_workers
self.prefetch_size = prefetch_size 预加载队列大小

self.data_queue = queue.Queue(maxsize=self.prefetch_size)
self.batch_indices = list(range(len(dataset)))
if self.shuffle:
random.shuffle(self.batch_indices)

self._stop_event = threading.Event()
self.worker_threads = []
for _ in range(self.num_workers):
thread = threading.Thread(target=self._worker, daemon=True)
self.worker_threads.append(thread)
thread.start()

self.current_batch_idx = 0
self.lock = threading.Lock()

def _worker(self):
while not self._stop_event.is_set():
try:
从队列中取一个 batch 的起始索引
batch_start_idx = self.current_batch_idx
if batch_start_idx + self.batch_size > len(self.dataset):
如果是最后一个 batch，并且不足一个 batch_size，也处理
batch_end_idx = len(self.dataset)
else:
batch_end_idx = batch_start_idx + self.batch_size

indices = self.batch_indices[batch_start_idx:batch_end_idx]
if not indices: 如果没有更多索引了
break

假装在这里进行复杂的数据加载和预处理
batch_data = []
for i in indices:
item = self.dataset[i] 实际从 dataset 加载
可以在这里进行更复杂的数据增强
batch_data.append(item)

将 batch 数据放入队列
self.data_queue.put(batch_data)

更新下一个 batch 的起始索引
with self.lock:
self.current_batch_idx += self.batch_size

except Exception as e:
print(f"Worker error: {e}")
break 出现错误时停止该 worker

def __iter__(self):
return self

def __next__(self):
if self.current_batch_idx >= len(self.dataset) and self.data_queue.empty():
所有数据都已处理完，重置并打乱索引，如果需要多轮训练
if self.shuffle:
random.shuffle(self.batch_indices)
self.current_batch_idx = 0
如果真的没有数据了，抛出 StopIteration
if self.data_queue.empty():
raise StopIteration

从队列中获取一个 batch
return self.data_queue.get()

def __len__(self):
return len(self.dataset) // self.batch_size + (len(self.dataset) % self.batch_size != 0)

def stop(self):
self._stop_event.set()
for thread in self.worker_threads:
thread.join()

使用示例（注意：这里的 dataset 只是一个简单的示例，实际你需要一个 PyTorch Dataset）
假设有一个 my_actual_dataset
dataloader = BackgroundDataLoader(my_actual_dataset, batch_size=32, shuffle=True, num_workers=4, prefetch_size=5)
for batch in dataloader:
训练代码...
dataloader.stop() 训练结束后调用，清理线程
```
重要提示： 上面那个 `BackgroundDataLoader` 是一个概念性示例，它省略了 `pin_memory`、`collate_fn` 等很多 `DataLoader` 的复杂功能。实际应用中，你需要更精细地处理这些细节，或者考虑使用 `torch.utils.data.DataLoader` 的更高级配置。

第五招：优化你的 `__getitem__` 和 `collate_fn`

别光顾着调 DataLoader 参数，你自己的数据处理逻辑才是根本。

`__getitem__` 里的 IO 和计算：
避免重复 IO： 如果同一个文件被多个 batch 访问，考虑是否能缓存。
文件格式： 某些文件格式（如 HDF5, TFRecord）可能比大量小文件（如 JPG）更适合批量读取，尤其是在有 `num_workers` 的情况下。
数据增强：
GPU 加速： 对于图像数据，像 `torchvision.transforms` 里的很多操作（如 RandomResizedCrop, RandomHorizontalFlip）最终会在 CPU 上执行。如果你的增强很复杂，可以考虑使用 GPU 上的数据增强库，如 Albumentations (它有 GPU 加速选项) 或者 DALI (NVIDIA Data Loading Library)，后者是专门为 GPU 加载和增强数据设计的，效果非常震撼，但学习曲线也比较陡峭。
延迟增强： 只在必要时进行复杂的数据增强，或者只对需要增强的数据进行。
`collate_fn` 的效率：
默认 `collate_fn`： PyTorch 的默认 `collate_fn` 已经做了很多优化。
自定义 `collate_fn`： 如果你自定义了 `collate_fn`，确保里面的操作是高效的。例如，如果你需要拼接多个 tensor，可以考虑 `torch.cat()`，而不是手动循环。如果你的数据类型不一致，确保转换是高效的。
数据类型： 尽量使用 `float32` 或者 `float16` (如果你的 GPU 和模型支持)。避免使用 `float64`，它会显著增加内存占用和计算量。

第六招：使用 DALI NVIDIA Data Loading Library (进阶)

如果你已经把上面能调的参数都调了，CPU 依然是瓶颈，并且你用的是 NVIDIA GPU，那么 DALI 绝对值得尝试。

怎么回事？ DALI 是 NVIDIA 开发的一个高性能数据加载和预处理库，它将数据加载和预处理的整个流程都放在 GPU 上执行（或利用 GPU 加速 CPU 操作）。这意味着你的 CPU 可以完全解放出来，专注于模型训练。
为什么强大？
GPU 加速： 图像解码、resize、crop、color jitter 等几乎所有常见的数据增强操作，DALI 都能在 GPU 上完成。
流水线优化： DALI 构建了一个高效的数据加载和处理流水线，可以实现端到端的优化。
格式支持： 支持 JPEG, PNG, TIFF 等多种图像格式，以及 TFRecord, MXNet RecordIO 等。
怎么用？
安装： 首先需要安装 DALI，通常通过 pip：`pip install nvidiadali`。
编写 DALI 管道： 你需要用 Python 编写 DALI 的操作来定义数据加载和预处理流程。
集成到 PyTorch： DALI 提供与 PyTorch 的集成接口，你可以创建一个 DALI 的 `DataLoader`。
```python
示例（简化版，实际 DALI 代码更复杂）
import nvidia.dali.fn as fn
import nvidia.dali.types as types
from nvidia.dali.pipeline import Pipeline
from nvidia.dali.plugin.pytorch import DALIClassificationIterator

class HybridPipe(Pipeline):
def __init__(self, batch_size, num_threads, device_id, data_dir, crop_size=224):
super(HybridPipe, self).__init__(batch_size, num_threads, device_id, seed=42)
self.input = fn.readers.tfrecord(
data_dir,
shard_files=True,
shard_idx=0, 根据你的多进程设置调整
read_data=True,
cycle=True,
pad_last_batch=True,
num_shards=1 根据你的多进程设置调整
)
DALI 的操作都在这里定义，使用 GPU 加速
self.decode = fn.decoders.image(self.input, output_type=types.RGB)
self.cmn = fn.crop_mirror_normalize(
self.decode,
... 各种增强参数 ...
)

def define_graph(self):
return self.cmn

训练时
pipe = HybridPipe(batch_size=32, num_threads=4, device_id=0, data_dir="/path/to/tfrecords")
pipe.build() 必须 build

DALI 提供 PyTorch Iterator
这个 iterator 实际上是 PyTorch DataLoader 的替代品，直接输出 GPU tensor
你可以直接在训练循环中使用它，无需再包装成 PyTorch DataLoader
train_loader = DALIClassificationIterator(pipe, reader_name="cuda_reader")

for batch in train_loader:
inputs, labels = batch[0]["data"], batch[0]["label"] DALI iterator 返回的结构
训练...
```
权衡：
学习曲线： DALI 的 API 和概念需要时间去理解，尤其是流水线和 GPU 操作的组合。
数据格式： DALI 对某些数据格式（如 TFRecord）的支持更友好，如果你的数据不是这些格式，可能需要先转换。
硬件要求： DALI 依赖 NVIDIA GPU。

总结一下我的经验：

1. 必选项： `num_workers` 和 `pin_memory=True` 是基础中的基础，几乎所有情况都应该先启用。
2. 微调： `prefetch_factor` 可以用来进一步压榨性能，但要小心内存。
3. 检查点： 仔细审视你的 `__getitem__` 和 `collate_fn`，看看有没有可以优化的 I/O 或计算。
4. 数据增强： 如果数据增强很复杂，考虑 GPU 加速或 DALI。
5. 极端情况： 对于超大规模数据集或极度复杂的预处理，DALI 是最后的杀手锏，但需要投入更多精力去学习和集成。

一些调试技巧：

隔离问题： 先尝试用最简单的数据集和最少的预处理来测试 DataLoader，看看速度有没有改善。如果连这个都慢，那问题可能出在更底层。
逐步增加复杂度： 每次只修改一个参数或一个优化方法，然后观察效果。避免一次性改太多，那样很难定位问题。
使用 `torch.profiler`： PyTorch 提供了强大的 profiler 工具，可以帮你分析训练过程中哪些部分耗时最多，包括数据加载。
查看日志： 很多时候，DataLoader 的问题会在日志中留下线索，比如内存警告、死锁提示等。

希望这些经验能给大家一些启发，让大家摆脱 DataLoader 慢的困扰！训练愉快！

如果可以的话，换硬盘。我猜提问者的数据是放在了机械硬盘上吧？其实你可以测试一下是不是单纯的数据读取阶段特别慢，因为根据我的经验这么长的读取时间已经不是transform造成的了。测试方法也很简单，你写dataset的时候，不要用opencv或者pillow对图片解析，直接用open()和read()读取每个图片的二进制数据（我没记错的话其实pillow中open图片时就是这么做的，你可以看看pillow中的源代码），dataloader返回的数据只是图片文件的二进制数据，这样测试dataloader，看看每一批时间是多久；可以读取每一批数据之后做一个delay，模拟训练模型花费时间。如果这样测试后发现每一批或者每隔几批后时间很久，最好是换硬盘吧。另外还有个简单的测试方法：移动/复制数据。比如把数据集的一部分或者找几个大文件，从硬盘复制到/dev/shm中，其实相当于从硬盘复制到内存中，计一下时间，就能得到传输速度，也可以根据你复制的文件量或者文件大小，换算出复制的数据集百分比，进而估计出一批数据读入内存所需的时间。

我在pytorch的数据读取上花了很大工夫做尝试，我说一下我的经历和经验。

我是做视频的，先说一下我的困境。假设，每一个batch包含16个视频，要在每个视频中取8帧，也就是送入模型的一个batch总共有128张图片。读取这128张图片有两种方法，一种方法是先读取整个mp4格式的视频，用openCV对每一帧解码然后返回要使用的帧，或用opencv里video.set方法跳跃到某帧对应的index只对该帧解码，当然前者运算量太大一般不用，使用后一种方式；另一种方法是提前将所有视频解码到图片，用的时候直接根据文件名读取图片就可以。这两种方法相比，第一种方法解码运算量大，因为从视频中解码帧比图片解码计算量大得多，查看CPU负载就可以看出差别来；第二种方法虽然快、计算量小，但是解码后的图片需要大量存储空间，保存到jpeg格式95%质量的情况下，像UCF101这样的视频数据集，图片数据是视频数据的6.5倍左右，Kinetics-400数据集解压到图片有1.5TB左右，如果想要保存成bmp格式或者numpy矩阵更不可能（一个图片的bmp格式与npy文件大小相同）。两种方法都有非常致命的缺点，只不过是牺牲计算还是牺牲存储空间的区别了。一般来说还是牺牲存储空间更合适，因为解码视频非常慢，即使数据从硬盘上的读取速度快，也会造成显卡有大量空余时间，并且CPU一直高负载运行也不利于服务器做其他任务（ssh和vim都卡）。另外，openCV从avi/mp4等格式里解析出来的图片数与用get()方法得到的视频帧数不一定相同（原因网上有解释），会导致帧检索错误。因此，我主要针对读取图片的方法尝试做改进。

1. 我一开始使用pre_fetch（可以在其他回答中看到类似的方法）办法，但是并没有用。
2. 我使用DALI做视频预处理，但是DALI的视频处理当时(2019.09)并不完善，我发现了bug……就反馈给开发者了，回复我说之后的版本会修复并做了标记。后来我也没用过DALI就没再看过了。
3. 我看到一些迷惑人的博客或帖子，以为pytorch的dataloader多进程是导致变慢的元凶，所以去研究了dataloader的源码。Dataloader工作时，首先主进程取一批数据在数据集中的index，然后其他进程（进程数由numworkers控制）中的一个进程对这一批数据进行读取和处理，处理完的数据放到dataqueue中，主进程从dataqueue中读取数据，然后再给一批数据的index。在dataloader开始使用时，都会出现数据加载时间很长的情况，这是因为第一次numworkers个进程都要产生2批数据。每次主进程产生数据index的代价都特别小，其实就是选择几个随机数字，网上有的地方说这一步导致数据不能并行，其实是错误的。但是，每个进程读取一批数据的时候是串行的，这一批中的数据一个一个读取。我就针对这一步做了改进尝试，在这个串行读取一批数据的代码里，源码有单线程限制语句，我就去掉了改成了多线程，我对这方面了解不够专业，不知道python里多进程加上多线程对不对，反正就试试吧。结果是，并没有提速效果。
4. 后来我专门测试了数据读取和数据转换的时间，发现经常会读取一张图片特别慢，1秒到2秒才能读到，几乎隔几张图片就会有这种卡顿，我试了PIL和openCV还有其他方法，都没办法提高速度。于是我按照文章开头的方法，做了硬盘读取速度的测试，发现是硬盘读取速度的限制造成了数据读取慢，而图片transform时间很短，也一直很稳定。
5. 针对硬盘读取数据，我想了一些办法。硬盘每次读取数据的时候，如果数据是连续的，那么读取速度会快很多，如果数据是很多小文件，读取速度就很慢，这一点平时复制文件的时候也会碰到。一批数据其实才16个视频，如果一次读取整个视频，就能减少读取次数，但是读取视频再解码运算量大。在一些框架（caffee、DALI）中会用到LMDB和HDF5文件格式，我可以借助这两种格式来把整个数据集所有图片保存到一个文件，从而减小读取次数。我看到一个帖子通过实验发现h5格式更快，所以就决定使用h5格式。但如果把每张图片用openCV读取后保存到numpy矩阵再存入h5文件，占用的存储空间太大（前边介绍过原因），所以就干脆保存每一个图片的二进制数据到numpy矩阵再存入h5文件，所有图片都读入，索引就是相对路径，这样做好了数据集文件。在使用的时候却发现了问题，h5文件不支持多进程读取，numworkers不能大于1，否则会报错；我又尝试了另一种格式，不是LMDB，而是npz格式，方法类似h5文件，但同样不支持多进程读取。
6. 单个h5和npz文件不支持，我就按视频拆分，每个视频保存一个h5文件，每个h5文件保存的是该视频每张图片的二进制数据，这样读取1个视频的时候一次性读取1个h5文件，取出里边的8张图片再解码就可以，相比直接读取视频找到帧再解码，这种解码图片的方法运算量低很多，当然对于同一个视频h5文件要比mp4/avi文件大很多，所以这种方法算是对运算量和占用空间的平衡吧。这种方法当时把训练时间降低了一点，效果不大。
7. 根据我的测试，h5文件读取数据的时候也并不是一次性读取整个文件的数据的，而是根据索引（或者说key）读取检索的数据，只不过一个h5文件是连续存储的。我还是想要一次性读取整个视频的内容，经过我测试，npy文件每次都是返回整个文件的数据，那我可以把每个视频所有图片的数据存在npy文件里。可是不同的图片数据长度不一样，一个npy文件只能保存维度相同的向量，所以我再一次牺牲空间，每个npy文件保存的矩阵里，矩阵的维度一个是视频帧数，另一个是这个视频所有图片中文件数据最长的长度，其他长度不足的图片后续补0。这样可以读入整个视频文件，直接根据矩阵索引找到对应的图片，读入图片未解码的数据后，通过PIL解码就可以还原图片了。当时我使用这个方法，把训练时间（间隔不是1个batch）从18秒左右减小到13秒左右。所以我一直使用这种方法了。
8. 但是我还是不满意。我想能不能直接从内存中读取npy文件呢？一次性加载所有数据集是不可能的，内存没那么大，那我可不可以一边把硬盘上的数据复制到内存里，一边读取内存的数据呢？可以先提前复制一些数据，给硬盘留一些时间余地。内存位置可以用/dev/shm，让dataloader到这里来读数据。我写程序做了试验，结果是这个方法不管用，因为从硬盘复制到/dev/shm太慢了。这一步也让我认识到了，数据加载慢只能怪硬盘。
9. 没有别的办法了，现在我就是使用7里边的npy文件数据集或者原始的视频文件，数据集不大的情况下会先把数据集复制到主硬盘SSD里，这个方法挺爽的。
10. 还没有结束，最近看到一篇文章：

“数据回波”，简单来说就是重复使用已经训练过还在内存中的数据。我还没有做试验。

更新：

之前实验了“数据回波”，太影响精度，提上来的一两个点又被抹平了，所以不建议使用。数据速度要加载快，归根到底还是要换固态或者内存，另外不要两个人同时从硬盘上读，可以分硬盘用。

又更新：

有些回答是针对数据预处理速度慢的，如果真的是这一步慢，也没必要使用DALI什么的，torchvision的Transforms中，几乎所有操作都支持PIL.Image和torch.Tensor两种输入，因此将图片读入后，先转换成Tensor（甚至torchvision有个read_image函数可以直接返回Tensor），然后放到GPU上，再用Transforms预处理，也是可以的。这个办法我已经用过了，在一个任务中我发现数据增强里的色彩增强最影响速度，其次是旋转，其他resize、flip都很快，因此用上述方法，在GPU上预处理数据。弊端是需要占用不少显存，我当时用的V100，还倒吃得消。为了能再节省显存，我把数据增强前的Tensor转换成了float16，送入网络之前再转换成float32，中间处理过程中，某些计算函数可能不支持float16，因此需要转换到float32再转换回float16。还需要注意的是，为了避免所有数据都只在GPU 0上预处理，要对不同的数据处理进程使用不同的GPU，比如可以用pid%n_gpu来指定。虽然麻烦点，哪怕中间还有float16和float32的来回转换，但是不怎么影响速度，数据处理的显存使用也能降低一些。亲测好用。