神经网络训练中高效读写图像与存储DataFrame的方法咨询

针对你遇到的18万+PNG图片读取内存占用高、Colab训练崩溃的问题，我从高效数据读取和DataFrame存储优化两个方向给出具体方案：

一、更高效的数据读取方法 #

核心思路是避免一次性加载所有图片到内存，采用**按需加载（懒加载）**的方式，训练时仅读取当前批次的图片，大幅降低内存压力。

1. 使用深度学习框架的Dataset/DataLoader（推荐） #

PyTorch或TensorFlow都提供了专门的数据集加载工具，能自动处理批次读取、打乱、多线程加载等，非常适合大规模数据集：

PyTorch示例：

import torch
from torch.utils.data import Dataset, DataLoader
import cv2
import glob
from torchvision import transforms

# 自定义数据集类，仅在需要时读取图片
class CustomImageDataset(Dataset):
    def __init__(self, image_paths, transform=None):
        self.image_paths = image_paths
        self.transform = transform  # 预处理操作

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

    def __getitem__(self, idx):
        # 仅在获取该样本时读取图片
        img_path = self.image_paths[idx]
        # 可以指定读取模式减少内存，比如灰度图用cv2.IMREAD_GRAYSCALE
        image = cv2.imread(img_path, cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转RGB（如果需要）
        
        # 应用预处理
        if self.transform:
            image = self.transform(image)
        
        return image  # 如果有标签，也可以一起返回

# 获取所有图片路径（这一步内存占用极低）
folders = glob.glob(r'path\to\images\*')
imagenames_list = []
for folder in folders:
    imagenames_list.extend(glob.glob(folder + '/*.png'))

# 定义预处理流程（按需调整）
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 创建数据集和DataLoader，设置合适的batch_size
dataset = CustomImageDataset(imagenames_list, transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

# 训练时迭代dataloader即可，每次仅加载一个批次的图片
for batch_images in dataloader:
    # 在这里执行训练步骤
    pass

TensorFlow示例：

import tensorflow as tf
import glob

# 获取所有图片路径
folders = glob.glob(r'path\to\images\*')
imagenames_list = []
for folder in folders:
    imagenames_list.extend(glob.glob(folder + '/*.png'))

# 构建tf.data.Dataset
dataset = tf.data.Dataset.from_tensor_slices(imagenames_list)

# 定义读取和预处理函数
def load_and_preprocess_image(path):
    image = tf.io.read_file(path)
    image = tf.image.decode_png(image, channels=3)
    image = tf.image.resize(image, (224, 224))
    image = tf.cast(image, tf.float32) / 255.0  # 归一化
    return image

# 应用映射和批次处理
dataset = dataset.map(load_and_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.shuffle(buffer_size=1000).batch(32).prefetch(tf.data.AUTOTUNE)

# 训练时迭代dataset
for batch_images in dataset:
    # 执行训练步骤
    pass

2. 优化单张图片读取 #

• 读取时指定图片格式：如果是灰度图，用cv2.imread(path, cv2.IMREAD_GRAYSCALE)，减少通道数（从3通道变为1通道），内存直接减少2/3。
• 使用PIL的懒加载：PIL.Image.open(path)不会立即把图片读进内存，直到调用image.load()或进行预处理操作，结合torchvision.transforms非常灵活。

二、更高效的DataFrame存储方法 #

首先要明确：DataFrame并不适合直接存储大量图片数组，因为图片是二进制数据，结构化存储会带来额外开销。更优的思路是：

1. 存储图片路径而非图片数据（推荐） #

把图片路径存入DataFrame，训练时通过路径读取图片，这样DataFrame体积极小，读写速度极快：

import pandas as pd

# 构建仅包含路径的DataFrame，可额外添加标签、预处理参数等列
df = pd.DataFrame({
    'image_path': imagenames_list,
    # 'label': your_labels_list  # 如果有标签的话
})

# 保存为pickle或更高效的格式
df.to_pickle(r'data\image_metadata.pkl')

# 读取后结合前面的Dataset使用
df_loaded = pd.read_pickle(r'data\image_metadata.pkl')
image_paths = df_loaded['image_path'].tolist()

2. 若必须存储图片数据，更换存储格式 #

如果业务需求必须把图片数组存在DataFrame里，替换pickle为以下更高效的格式：

Feather格式

专为Pandas DataFrame设计的二进制格式，读写速度快，占用空间远小于pickle：

# 需要先安装pyarrow：!pip install pyarrow
df.to_feather(r'data\images_data.feather')
df_loaded = pd.read_feather(r'data\images_data.feather')

Parquet格式

列存储格式，支持高效压缩，适合大规模数据集：

# 安装依赖：!pip install pyarrow
df.to_parquet(r'data\images_data.parquet', compression='snappy')
df_loaded = pd.read_parquet(r'data\images_data.parquet')

HDF5格式

支持分块存储和压缩，可按需加载部分数据，适合超大数据集：

df.to_hdf(r'data\images_data.h5', key='images', mode='w', complevel=5, complib='blosc')
# 读取时可指定行列，减少内存占用
df_loaded = pd.read_hdf(r'data\images_data.h5', key='images', where='index < 1000')

3. 单独存储图片数组 #

如果不需要DataFrame的结构化特性，可以把预处理后的图片数组用numpy压缩存储：

import numpy as np

# 假设read_images是预处理后的numpy数组列表（需保证所有图片形状一致）
images_array = np.array(read_images)
# 压缩存储
np.savez_compressed(r'data\images_compressed.npz', images=images_array)

# 读取
loaded_data = np.load(r'data\images_compressed.npz')
images_loaded = loaded_data['images']

额外Colab优化建议 #

• 挂载Google Drive：把数据集存到Drive，避免每次重新上传，同时可直接读取Drive中的数据。
• 启用GPU/TPU：Colab的GPU/TPU不仅加速训练，还会分配更多内存资源。
• 定期清理内存：训练间隙调用import gc; gc.collect()释放未使用的内存。

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