如何一次性预处理图像数据集，实现多轮模型训练测试复用？

完全可以一次性完成预处理并存储，后续直接加载已处理好的数据，避免重复计算，尤其适合多轮实验场景。下面是具体实现步骤：

一、存储预处理后的数据集 #

由于归一化后的Tensor可能包含负数，无法直接保存为JPG/PNG等常规图像格式，我们可以将预处理后的Tensor以PyTorch二进制格式（.pt）保存，同时保留原数据集的类别文件夹结构，方便后续加载。

代码实现 #

import os
import torch
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder

# 原数据集路径
TRAIN_ROOT = "你的训练集根路径"
TEST_ROOT = "你的测试集根路径"
# 预处理后数据的存储路径
SAVED_TRAIN_ROOT = "预处理后训练集存储路径"
SAVED_TEST_ROOT = "预处理后测试集存储路径"

# 定义训练集和测试集的预处理逻辑
train_transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.541, 0.536, 0.357],
                         std=[0.321, 0.339, 0.441])
])

test_transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.540, 0.536, 0.357],
                         std=[0.321, 0.339, 0.441])
])

# 批量处理并保存数据的函数
def save_preprocessed_data(raw_root, save_root, transform):
    os.makedirs(save_root, exist_ok=True)
    dataset = ImageFolder(root=raw_root, transform=transform)
    
    for idx, (tensor_img, label) in enumerate(dataset):
        # 获取原图像的相对路径，保持类别文件夹结构
        img_path, _ = dataset.samples[idx]
        relative_path = os.path.relpath(img_path, raw_root)
        # 替换文件后缀为.pt
        save_path = os.path.join(save_root, os.path.splitext(relative_path)[0] + ".pt")
        # 创建类别子文件夹
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        # 保存预处理后的Tensor
        torch.save(tensor_img, save_path)
        
        # 打印进度（可选）
        if (idx+1) % 1000 == 0:
            print(f"已处理 {idx+1}/{len(dataset)} 张图像")

# 执行训练集和测试集的预处理存储
save_preprocessed_data(TRAIN_ROOT, SAVED_TRAIN_ROOT, train_transform)
save_preprocessed_data(TEST_ROOT, SAVED_TEST_ROOT, test_transform)

二、加载预处理后的数据集 #

自定义一个Dataset类，直接读取已保存的.pt文件，训练时无需再执行任何预处理步骤。

代码实现 #

import os
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.datasets.folder import find_classes, make_dataset

class PreprocessedImageDataset(Dataset):
    def __init__(self, root):
        # 获取类别列表和类别索引映射
        self.classes, self.class_to_idx = find_classes(root)
        # 获取所有.pt格式的样本路径
        self.samples = make_dataset(root, self.class_to_idx, extensions=(".pt",))
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        pt_path, label = self.samples[idx]
        # 加载预处理好的Tensor
        tensor_img = torch.load(pt_path)
        return tensor_img, label

# 加载预处理后的数据集
train_dataset = PreprocessedImageDataset(SAVED_TRAIN_ROOT)
test_dataset = PreprocessedImageDataset(SAVED_TEST_ROOT)

# 创建DataLoader用于训练
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

注意事项 #

• 磁盘空间：30000张224×224的3通道Tensor，单张约600KB，总容量约18GB，需提前确认磁盘剩余空间。
• 参数变更：若后续需要调整预处理参数（如Resize尺寸、归一化均值/标准差），需重新执行存储步骤。
• 加载效率：设置num_workers时需结合磁盘IO性能调整，避免出现加载瓶颈。

内容的提问来源于stack exchange，提问作者Prajakta Rathod