训练深度学习模型时能否预转换图像？预存张量方案效率探讨

在训练深度学习模型时，PyTorch的DataLoader有时会因为实时执行数据变换成为性能瓶颈。以torchvision.datasets中DatasetFolder的__getitem__函数为例：

path, target = self.samples[index]
sample = self.loader(path)
if self.transform is not None:
    sample = self.transform(sample)
if self.target_transform is not None:
    target = self.target_transform(target)

return sample, target

我想知道是否可以提前将ImageNet这类数据集的图像预处理为张量并保存到磁盘，然后修改__getitem__函数直接从磁盘读取这些张量？这种方法的效率如何？有没有人尝试过这个方案？

我担心从磁盘加载会成为新的瓶颈，替代之前的数据变换瓶颈。另外还有存储空间问题：比如用以下标准变换处理后，单张ImageNet图像保存为张量大小可达74MB：

transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])

有没有人试过这个方案？当然有 #

这个思路完全可行，而且在ImageNet这类大规模数据集的训练场景中已经被广泛实践。尤其是当你需要反复训练（比如调参、复现实验），提前做一次预处理能彻底避免每次训练都重复跑相同的变换逻辑，省下来的CPU资源可以用来做其他事，或者配合GPU更高效地并行。

效率到底怎么样？看这两个关键因素 #

• 如果数据变换是CPU瓶颈：如果你的变换包含大量CPU密集型操作（比如复杂的图像增强、多轮计算型预处理），预存张量绝对能提速——把这部分计算从训练流程里剥离，让CPU不用在训练时实时跑变换，转而专注于数据加载的IO调度，此时磁盘读取的开销通常远小于实时变换的CPU开销，整体训练速度会明显提升。
• 如果磁盘IO拖后腿：要是用的是普通机械硬盘（HDD），大量小文件的随机读取确实可能变成新瓶颈，但换成SSD尤其是NVMe SSD的话，这个问题基本不存在。另外还可以用批量存储优化：比如把几十上百个张量打包成一个HDF5、LMDB文件，减少文件数量，提升读取效率。

存储空间问题怎么解决？这几招管用 #

你提到的单张74MB是未压缩的32位浮点数张量，其实有很多办法大幅压缩存储占用：

• 无损压缩存储：用numpy的npz格式、或者PyTorch保存时开启compress=True，单张张量能压缩到10-20MB，完全不损失精度。
• 低精度存储：如果训练支持半精度（FP16），把张量转成torch.float16再存储，存储空间直接减半，现在大部分现代GPU都支持FP16训练，对模型精度影响极小。
• 用专门的存储容器：LMDB、HDF5这类格式不仅能减少文件系统的管理开销，还支持高效的批量读取，自带的压缩功能也能进一步节省空间。

实操建议 #

• 编写预处理脚本：遍历整个数据集，执行所有固定变换（比如Resize、CenterCrop、ToTensor、Normalize），将处理好的张量和对应标签按类别或批量分组存储（比如每个类别一个文件夹，或每1000张存一个打包文件）。
• 自定义Dataset类：修改__getitem__函数，直接从磁盘读取预存的张量和标签，无需再执行变换操作。
• 优化DataLoader配置：开启num_workers多进程加载，同时设置pin_memory=True（使用GPU训练时），进一步加快数据从内存到GPU的传输速度。

内容的提问来源于stack exchange，提问作者Giang Nguyen