如何将含二维数组的pandas DataFrame转换为torch Dataset用于ASR微调

1. 将pandas DataFrame转换为符合要求的Dataset类型 #

你遇到的Dataset.from_pandas报错是因为speech、input_values列每个单元格存储的是嵌套的二维列表，处理后即可正常转换，有两种常用实现方案：

方案一：自定义PyTorch Dataset（最灵活，适配任意数据格式） #

直接继承torch.utils.data.Dataset实现自定义数据集类，无需修改原有DataFrame的存储结构：

from torch.utils.data import Dataset

class ASRDataset(Dataset):
    def __init__(self, df):
        # 重置索引避免拆分数据集后索引不连续导致的取值错误
        self.df = df.reset_index(drop=True)
    
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        return {
            # 去掉外层嵌套的空列表维度，拿到实际的特征数组
            "input_values": row["input_values"][0],
            "labels": row["labels"],
            # 如有需要可以新增其他字段，比如转写文本用于验证指标
            "sentence": row["sentence"]
        }

实例化直接传入拆分后的DataFrame即可：

ds_train = ASRDataset(train_df)
ds_test = ASRDataset(test_df)

方案二：预处理后使用Hugging Face Dataset转换 #

先将二维嵌套列拍平为一维，再调用from_pandas转换：

from datasets import Dataset

# 定义预处理函数，去掉嵌套的外层列表维度
def flatten_nested(row):
    row["speech"] = row["speech"][0]
    row["input_values"] = row["input_values"][0]
    return row

# 处理数据集后转换
ds_train = Dataset.from_pandas(ds_train.apply(flatten_nested, axis=1))
ds_test = Dataset.from_pandas(ds_test.apply(flatten_nested, axis=1))

# 可选：设置数据集格式为torch，返回的字段直接为Tensor类型，无需额外转换
ds_train.set_format("torch", columns=["input_values", "labels"])
ds_test.set_format("torch", columns=["input_values", "labels"])

2. 训练验证阶段data_collator调用说明 #

是的，训练和验证阶段都会调用传入的data_collator。两个阶段都需要将单个样本组合为批次、完成长度对齐、掩码生成等批处理操作，所以都会走data_collator的处理逻辑。

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