使用Keras训练UNet时大内存数据分批加载问题求助

问题原因 #

Keras无法直接识别RaggedMmap类型的数据结构，需要将其转换为Keras兼容的数据加载方式，同时实现分批加载以节省内存。

优先方案：Keras兼容的分批加载实现 #

方案1：用tf.data.Dataset包装RaggedMmap #

直接将RaggedMmap的迭代逻辑包装成TensorFlow数据集，Keras可直接识别，自动分批加载并释放内存：

import tensorflow as tf
from mmap_ninja.ragged import RaggedMmap

# 加载内存映射数据
X_train_mmp = RaggedMmap('/kaggle/working/xTreino_mmap')
y_train_cat_mmp = RaggedMmap('/kaggle/working/yTreino_mmap')
X_test_mmp = RaggedMmap('/kaggle/working/xTeste_mmap')
y_test_cat_mmp = RaggedMmap('/kaggle/working/yTeste_mmap')

# 定义训练/验证数据生成器
def train_gen():
    for idx in range(len(X_train_mmp)):
        yield X_train_mmp[idx], y_train_cat_mmp[idx]

def val_gen():
    for idx in range(len(X_test_mmp)):
        yield X_test_mmp[idx], y_test_cat_mmp[idx]

# 转换为tf.data.Dataset，指定输出张量签名
batch_size = 8
train_dataset = tf.data.Dataset.from_generator(
    train_gen,
    output_signature=(
        tf.TensorSpec(shape=X_train_mmp[0].shape, dtype=tf.float32),
        tf.TensorSpec(shape=y_train_cat_mmp[0].shape, dtype=tf.float32)  # 标签类型按需调整
    )
).batch(batch_size).prefetch(tf.data.AUTOTUNE)

val_dataset = tf.data.Dataset.from_generator(
    val_gen,
    output_signature=(
        tf.TensorSpec(shape=X_test_mmp[0].shape, dtype=tf.float32),
        tf.TensorSpec(shape=y_test_cat_mmp[0].shape, dtype=tf.float32)
    )
).batch(batch_size).prefetch(tf.data.AUTOTUNE)

# 启动训练
history = model.fit(
    train_dataset,
    verbose=1,
    epochs=10,
    validation_data=val_dataset,
    shuffle=False  # 如需洗牌，可在generator中添加逻辑或使用dataset.shuffle
)

方案2：自定义Sequence类（更灵活） #

如果需要更精细的加载控制（比如直接操作bin文件指针），继承Keras的Sequence类实现自定义分批加载：

from keras.utils import Sequence
import numpy as np

class BinDataLoader(Sequence):
    def __init__(self, x_bin_path, y_bin_path, batch_size, total_samples, sample_shape, dtype=np.float32):
        self.x_path = x_bin_path
        self.y_path = y_bin_path
        self.batch_size = batch_size
        self.total_samples = total_samples
        self.sample_shape = sample_shape
        self.dtype = dtype
        self.sample_bytes = np.prod(sample_shape) * dtype.itemsize
        self.indices = np.arange(total_samples)

    def __len__(self):
        # 计算总批次数
        return int(np.ceil(self.total_samples / self.batch_size))

    def __getitem__(self, idx):
        # 加载当前批次样本
        batch_idx = self.indices[idx*self.batch_size : (idx+1)*self.batch_size]
        batch_x, batch_y = [], []
        
        with open(self.x_path, 'rb') as f_x, open(self.y_path, 'rb') as f_y:
            for i in batch_idx:
                # 定位到样本位置并读取
                f_x.seek(i * self.sample_bytes)
                x = np.frombuffer(f_x.read(self.sample_bytes), dtype=self.dtype).reshape(self.sample_shape)
                batch_x.append(x)
                
                f_y.seek(i * self.sample_bytes)  # 标签字节数不同需调整
                y = np.frombuffer(f_y.read(self.sample_bytes), dtype=self.dtype).reshape(self.sample_shape)
                batch_y.append(y)
        
        return np.array(batch_x), np.array(batch_y)

    def on_epoch_end(self):
        # 每个epoch后洗牌样本（可选）
        np.random.shuffle(self.indices)

# 初始化加载器
batch_size = 8
train_loader = BinDataLoader(
    x_bin_path='/kaggle/working/xTreino.bin',
    y_bin_path='/kaggle/working/yTreino.bin',
    batch_size=batch_size,
    total_samples=10000,  # 替换为实际训练样本数
    sample_shape=(256,256,3)  # 替换为你的样本形状
)

val_loader = BinDataLoader(
    x_bin_path='/kaggle/working/xTeste.bin',
    y_bin_path='/kaggle/working/yTeste.bin',
    batch_size=batch_size,
    total_samples=2000,  # 替换为实际验证样本数
    sample_shape=(256,256,3)
)

# 训练模型
history = model.fit(
    train_loader,
    verbose=1,
    epochs=10,
    validation_data=val_loader
)

备选方案：PyTorch分批加载 #

如果切换到PyTorch，使用Dataset+DataLoader实现分批加载：

import torch
from torch.utils.data import Dataset, DataLoader
from mmap_ninja.ragged import RaggedMmap

class MmapDataset(Dataset):
    def __init__(self, x_mmap, y_mmap):
        self.x_data = x_mmap
        self.y_data = y_mmap

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

    def __getitem__(self, idx):
        x = torch.tensor(self.x_data[idx], dtype=torch.float32)
        y = torch.tensor(self.y_data[idx], dtype=torch.float32)
        return x, y

# 加载数据
X_train_mmp = RaggedMmap('/kaggle/working/xTreino_mmap')
y_train_cat_mmp = RaggedMmap('/kaggle/working/yTreino_mmap')
X_test_mmp = RaggedMmap('/kaggle/working/xTeste_mmap')
y_test_cat_mmp = RaggedMmap('/kaggle/working/yTeste_mmap')

# 创建DataLoader
batch_size = 8
train_loader = DataLoader(MmapDataset(X_train_mmp, y_train_cat_mmp), batch_size=batch_size, shuffle=True)
val_loader = DataLoader(MmapDataset(X_test_mmp, y_test_cat_mmp), batch_size=batch_size)

# 训练循环示例
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters())
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(10):
    model.train()
    total_loss = 0.0
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        optimizer.zero_grad()
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            outputs = model(batch_x)
            val_loss += criterion(outputs, batch_y).item()
    
    print(f"Epoch {epoch+1} | Train Loss: {total_loss/len(train_loader):.4f} | Val Loss: {val_loss/len(val_loader):.4f}")

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