自编码器预训练Unet语义分割：数据集数量与内存优化问题

不需要强制使用相同数量的图像。

• 预训练自编码器的核心目标是学习通用的图像特征表示，只要预训练数据集和分割任务的图像是同分布的（比如都是颅骨图像），即使数据量少于U-Net的训练数据，也能学到有效的底层特征（如边缘、纹理），迁移到U-Net后依然能提升分割性能。
• 训练U-Net时，需要的是带标注的任务相关数据，数量取决于标注成本和任务复杂度：如果分割任务难度高（比如精细的颅骨结构分割），需要更多标注数据来学习细节；如果任务简单，少量标注数据配合预训练权重也能达到不错的IOU。
• 当然，如果两者数据量匹配且都是高质量标注数据，效果会更稳定，但这不是必须条件。关键是保证预训练数据和U-Net训练数据的分布一致，这样迁移的特征才有用。

原代码一次性将1600张512×512的图像加载到内存，每张float32格式的图像占约3MB（512×512×3×4字节），1600张总计约4.8GB，容易超出Colab实例的内存上限。解决思路是分批加载图像，不用一次性把所有数据存入内存，推荐两种实现方式：

方法1：使用TensorFlow Dataset API（推荐） #

TensorFlow的Dataset会按需加载和预处理图像，内存占用极低：

import tensorflow as tf
import os

SIZE = 512
path1 = '/content/drive/MyDrive/Colab Notebooks/Skull images'

# 获取所有图像文件路径（过滤非图像文件）
file_paths = [os.path.join(path1, fname) for fname in os.listdir(path1) if fname.endswith(('.png', '.jpg', '.jpeg'))]

def load_and_preprocess(file_path):
    # 读取并解码图像
    img = tf.io.read_file(file_path)
    img = tf.image.decode_jpeg(img, channels=3)  # 若为PNG格式，替换为decode_png
    # 调整尺寸
    img = tf.image.resize(img, [SIZE, SIZE])
    # 归一化到[0,1]
    img = tf.cast(img, tf.float32) / 255.0
    return img

# 创建数据集并并行预处理
dataset = tf.data.Dataset.from_tensor_slices(file_paths)
dataset = dataset.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
# 设置批次大小，可根据内存调整
dataset = dataset.batch(16).prefetch(tf.data.AUTOTUNE)

训练时直接将dataset传入model.fit()即可。

方法2：自定义Keras生成器 #

如果习惯用Keras的Sequence类，可实现一个分批加载的生成器：

import cv2
import numpy as np
import os
from tensorflow.keras.utils import Sequence

SIZE = 512
path1 = '/content/drive/MyDrive/Colab Notebooks/Skull images'
file_list = [fname for fname in os.listdir(path1) if fname.endswith(('.png', '.jpg', '.jpeg'))]
batch_size = 16

class ImageBatchGenerator(Sequence):
    def __init__(self, file_list, img_dir, img_size, batch_size):
        self.file_list = file_list
        self.img_dir = img_dir
        self.img_size = img_size
        self.batch_size = batch_size

    def __len__(self):
        # 返回总批次数
        return int(np.ceil(len(self.file_list) / self.batch_size))

    def __getitem__(self, idx):
        # 加载当前批次的图像
        batch_files = self.file_list[idx*self.batch_size : (idx+1)*self.batch_size]
        batch_imgs = []
        for fname in batch_files:
            img = cv2.imread(os.path.join(self.img_dir, fname), 1)
            img = cv2.resize(img, self.img_size)
            img = img.astype('float32') / 255.0
            batch_imgs.append(img)
        return np.array(batch_imgs)

# 初始化生成器
img_generator = ImageBatchGenerator(file_list, path1, (SIZE, SIZE), batch_size)

训练时使用model.fit(img_generator, epochs=...)即可。

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