本地大Parquet文件流式输入HuggingFace Trainer仅加载首个小批量问题

我尝试将本地大型.parquet文件流式输入到HuggingFace Trainer训练分类模型，但脚本只加载了第一个小批量数据——明明文件体积很大，1个epoch本该耗时约10小时，实际epoch数却快速增长。

问题根源 #

自定义MyIterableDataset的__iter__方法存在逻辑错误：

• 仅调用一次next(self.generator)，只获取了Parquet文件的第一个batch
• 返回的迭代器仅包含该batch的样本，Trainer会误认为整个数据集已遍历完成，直接进入下一个epoch

修正后的代码 #

修改MyIterableDataset，让它持续遍历所有Parquet batch并逐个返回样本：

import pyarrow.parquet as pq
import torch
import pandas as pd
import evaluate
from transformers import (AutoTokenizer, CamembertForSequenceClassification,
                          EarlyStoppingCallback, Trainer, TrainingArguments,
                          pipeline)
import numpy as np

class MyIterableDataset(torch.utils.data.IterableDataset):

    def __init__(self, parquet_file_path: str, tokenizer, label_encoder, batch_size: int = 8):
        self.parquet_file = pq.ParquetFile(parquet_file_path)
        self.batch_size = batch_size
        self.tokenizer = tokenizer
        self.label_encoder = label_encoder

    def __iter__(self):
        # 遍历Parquet文件的所有batch
        for batch in self.parquet_file.iter_batches(batch_size=self.batch_size):
            encodings = self.tokenizer(batch['text'].tolist(), truncation=True, padding=True, max_length=512)
            # 逐个生成样本
            for idx in range(len(batch)):
                item = {key: torch.tensor(val[idx]) for key, val in encodings.items()}
                item["labels"] = torch.tensor(self.label_encoder.transform([str(batch['target'][idx])]))
                yield item

关键改动说明 #

1. 使用yield关键字将__iter__转为生成器，持续返回样本直到所有batch遍历完成
2. 删除初始化时的self.generator，改为在__iter__中直接遍历iter_batches，避免生成器耗尽后无法重置的问题
3. 验证集的MyIterableDataset需要做同样修改，确保验证阶段能遍历所有数据

额外注意 #

• 若设置了max_steps=100000，训练会在达到10万步后停止，不受num_train_epochs限制，可根据需求调整该参数
• 确保label_encoder对目标值的转换逻辑稳定，避免出现编码错误

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