Python中如何压缩超出可用内存的列表并构建PyTorch图像数据集

你当前的实现是把所有图像预处理后的张量全部存入内存列表，再统一写出为单个大文件，26GB的张量数据在存储时会额外占用内存拷贝开销，很容易触发OOM（内存不足）导致进程终止。

不需要提前把所有数据加载到内存，仅在每个batch训练前动态加载、处理对应样本，内存占用仅和batch size正相关，完全可以承载26GB级别的数据集：

import os
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

class CustomImageDataset(Dataset):
    def __init__(self, root_dir, imsize=100):
        self.root_dir = root_dir
        self.imsize = imsize
        # 只预存所有样本的路径和标签，不加载图像本身，内存占用极低
        self.sample_paths = []
        self.labels = []
        # 遍历目录存储路径映射
        for label_idx, dir_name in enumerate(os.listdir(root_dir)):
            dir_path = os.path.join(root_dir, dir_name)
            if not os.path.isdir(dir_path):
                continue
            for file_name in os.listdir(dir_path):
                self.sample_paths.append(os.path.join(dir_path, file_name))
                self.labels.append(label_idx)
    
    def __len__(self):
        return len(self.sample_paths)
    
    def __getitem__(self, idx):
        img_path = self.sample_paths[idx]
        label = self.labels[idx]
        # 只有调用__getitem__时才加载处理单张图像
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        # 中心裁剪到1080x1080
        img = img[:, 420:1500]
        #  resize到目标尺寸
        img = cv2.resize(img, (self.imsize, self.imsize))
        # 归一化到[-1, 1]
        img = img.astype(np.float32) / 128.0 - 1.0
        # 转张量
        tensor = torch.from_numpy(img)
        return tensor, label

# 使用方式
dataset = CustomImageDataset(root="你的数据集根目录", imsize=100)
dataloader = DataLoader(
    dataset,
    batch_size=32, # 可根据你的显存容量调整
    shuffle=True, # 实现全局样本打乱
    num_workers=4, # 多进程加载，可按CPU核心数调整
    prefetch_factor=2, # 预取2个batch的数据，平衡加载速度和内存占用
    pin_memory=True # 用GPU训练时开启，可加速数据转移到GPU的速度
)

# 训练循环直接遍历dataloader即可
for batch_imgs, batch_labels in dataloader:
    # 此处写你的训练逻辑
    pass

• 本地缓存预处理结果：如果训练时IO/图像处理速度成为瓶颈，可以提前把每张图处理后的张量单独存为.pt文件到缓存目录，__getitem__优先读缓存文件，不需要重复做裁剪、resize操作，不需要占用多少内存就能大幅提升加载速度。
• 分片预存数据集：如果必须离线完成所有预处理，不要存成单个26GB的大文件，每1000~5000个样本存为一个小pt文件，再自定义Dataset按分片加载，用完即释放分片内存，也不会出现OOM。不需要手动拆分多个子集跑训练循环，会额外增加打乱逻辑的复杂度，也容易引入数据分布偏差。

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