如何为Huggingface Trainer创建加载多文件的自定义DataLoader

核心方案 #

由于内存限制无法合并所有数据集，我们可以采用分文件迭代+双层打乱的方式：每个epoch先随机打乱数据集文件的顺序，再对单个数据集开启内部shuffle加载batch。这样既满足"每个epoch文件顺序随机"，也保证"单个文件内数据顺序随机"，同时不会占用过多内存。

实现代码 #

import random
from torch.utils.data.dataloader import DataLoader

# 你的数据集列表，每个元素是包含input_ids的tokenized Dataset
tokenized_datasets = [tokenized_dataset_1, tokenized_dataset_2, ..., tokenized_dataset_10000]

# 训练循环
num_epochs = 5  # 替换成你的训练轮数
for epoch in range(num_epochs):
    # 1. 打乱当前epoch的数据集文件顺序
    random.shuffle(tokenized_datasets)
    
    # 遍历每个数据集文件
    for ds in tokenized_datasets:
        ds.set_format("torch")  # 转为PyTorch张量格式（如果未预先设置）
        # 单个数据集的DataLoader，shuffle=True保证文件内数据打乱
        train_loader = DataLoader(ds["train"], batch_size=32, shuffle=True)
        
        # 训练batch
        for batch in train_loader:
            # 因果LM的标签等于input_ids，直接构建输入
            model_inputs = {
                "input_ids": batch["input_ids"],
                "labels": batch["input_ids"]
            }
            outputs = model(**model_inputs)
            loss = outputs.loss
            
            # 反向传播与优化
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            
            # 按需添加日志、梯度累积、验证等逻辑

额外说明 #

• 如果所有数据集已经提前设置过set_format("torch")，可以把ds.set_format("torch")移到循环外，避免重复执行。
• 验证集的处理逻辑一致：创建验证数据集列表，每个epoch打乱后遍历每个验证集的DataLoader，进行评估即可。
• 这种方式每次仅加载单个数据集的batch到内存，完全规避合并数据集带来的内存压力。

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