Ubuntu下PyTorch数据加载怎样加速

Ubuntu下PyTorch数据加载加速方法

1. 多线程/多进程并行加载（核心优化）

通过DataLoader的num_workers参数启用多进程并行加载，将数据读取、预处理等CPU密集型操作分配给多个子进程，与GPU计算重叠执行，避免GPU等待数据。关键技巧：

数值设置：num_workers通常取CPU核心数或GPU数量的较小值（如4核CPU、1张GPU可设为4），避免过多进程导致资源竞争或内存溢出；
注意事项：设置后需测试不同数值的性能（如用time模块统计epoch时间），找到硬件适配的最优值。

2. 内存锁定（pin_memory）优化

启用pin_memory=True将数据锁定在CPU的连续物理内存区域，减少GPU访问时的内存拷贝开销（从“系统内存→GPU内存”的普通拷贝变为“锁定内存→GPU内存”的快速拷贝）。适用场景：GPU训练时，尤其对批量较大的数据集（如ImageNet）效果显著。

3. 数据预取（prefetch_factor）

通过prefetch_factor参数让DataLoader提前预取指定数量的批次数据（如prefetch_factor=2表示预取2个批次），减少I/O等待时间。示例代码：

dataloader = DataLoader(dataset, batch_size=32, num_workers=4, pin_memory=True, prefetch_factor=2)

注意：预取过多可能导致内存占用过高，需根据GPU显存和数据集大小调整。

4. 硬件级加速：使用RAMDisk存储数据

将数据复制到Ubuntu的RAMDisk（内存虚拟的磁盘，如/dev/shm）中，利用内存的高速读写（比SSD快10倍以上）加速数据加载。操作步骤：

查看当前RAMDisk大小：df -h /dev/shm；
修改/etc/fstab文件，添加tmpfs /dev/shm tmpfs defaults,size=4G 0 0（将大小设为4GB，可根据内存调整）；
重新挂载：sudo mount /dev/shm；
将数据集移动到/dev/shm目录下使用。

5. 优化数据预处理流程

使用高效库：用numpy、opencv替代Python原生循环（如用cv2.resize代替PIL的resize），利用向量化操作提升预处理速度；
简化预处理：减少不必要的操作（如多余的通道转换、重复的归一化），仅保留对模型训练必要的步骤；
缩减数据大小：若数据集中包含无用信息（如高分辨率图像的非关键区域），可通过裁剪、降采样等方式减少数据量。

6. 使用高效数据加载库（如DALI）

用NVIDIA的DALI库（深度学习数据加载库）替代PyTorch原生Dataset，DALI支持GPU加速的数据预处理（如图像解码、缩放、裁剪），大幅提升大规模数据集的加载效率。示例代码：

from nvidia.dali.pipeline import pipeline_def
import nvidia.dali.fn as fn
import nvidia.dali.types as types

@pipeline_def
def image_pipeline(file_root, batch_size, num_threads, device_id):
    images, labels = fn.readers.file(file_root=file_root, batch_size=batch_size, num_threads=num_threads, device_id=device_id)
    images = fn.decoders.image(images, device="mixed", output_type=types.RGB)
    images = fn.resize(images, resize_x=224, resize_y=224)
    return images, labels

# 创建pipeline
pipe = image_pipeline(file_root="/path/to/data", batch_size=32, num_threads=4, device_id=0)
pipe.build()

# 在训练循环中使用
for i in range(num_epochs):
    for images, labels in pipe:
        # 将数据转为PyTorch张量并送入模型
        images = images.as_cpu().numpy()
        labels = labels.as_cpu().numpy()
        # ... 训练代码 ...

注意：DALI需与PyTorch集成（通过torch.utils.data.DataLoader包装），适合大规模图像数据集。

7. 数据缓存策略

内存缓存：对预处理后的数据进行缓存（如用joblib.dump/joblib.load或torch.save/torch.load），避免重复预处理；
磁盘缓存：将预处理后的数据保存到SSD中，下次训练时直接加载，减少原始数据的读取时间。

8. 减少小文件数量

若数据集由大量小文件（如10万张以上的图片）组成，合并为TFRecord或LMDB格式，减少文件系统查找时间（小文件读取的I/O开销远大于大文件）。示例：

用tf.io.TFRecordWriter将图片和标签写入TFRecord文件；
用tf.data.TFRecordDataset读取TFRecord文件，配合map函数解析数据。