处理变长序列数据：自定义PyTorch数据集与DataLoader

题目

处理变长序列数据：自定义PyTorch数据集与DataLoader

信息

• 类型：问答
• 难度：⭐⭐

考点

Dataset类继承, 数据预处理, DataLoader参数配置, 自定义collate_fn, 变长序列处理

快速回答

处理变长序列数据的关键步骤：

• 继承torch.utils.data.Dataset实现自定义数据集类
• 在__getitem__中返回单个样本的元组（如(sequence, label, length)）
• 自定义collate_fn函数：
  • 按序列长度降序排序
  • 使用pad_sequence进行零填充
  • 重组数据和标签
• 创建DataLoader时设置collate_fn参数和batch_size

问题场景

在NLP或语音处理中，常遇到变长序列数据（如不同长度的句子或音频片段）。直接批量处理会导致以下问题：

• 无法堆叠成规则张量
• 填充不当造成计算资源浪费
• RNN/LSTM需要长度信息进行高效计算

解决方案

1. 自定义Dataset类

import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence

class VariableLengthDataset(Dataset):
    def __init__(self, sequences, labels):
        """
        参数:
        sequences: 列表的列表 [[1,2], [3,4,5], ...]
        labels: 对应标签列表
        """
        self.sequences = sequences
        self.labels = labels

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

    def __getitem__(self, idx):
        seq = torch.tensor(self.sequences[idx], dtype=torch.long)
        label = torch.tensor(self.labels[idx], dtype=torch.float)
        length = torch.tensor(len(seq), dtype=torch.long)
        return seq, label, length  # 返回三元组

2. 自定义collate_fn函数

def custom_collate(batch):
    """
    处理批次数据:
    1. 按序列长度降序排序
    2. 分离数据、标签和长度
    3. 填充序列
    """
    # 按长度降序排序
    batch.sort(key=lambda x: x[2], reverse=True)
    sequences, labels, lengths = zip(*batch)

    # 填充序列（右侧填充0）
    padded_seqs = pad_sequence(sequences, batch_first=True, padding_value=0)

    # 转换为张量
    labels = torch.stack(labels)
    lengths = torch.stack(lengths)

    return padded_seqs, labels, lengths

3. 创建DataLoader

# 示例数据
sequences = [[1,2,3], [4,5], [6,7,8,9]]
labels = [0, 1, 0]

dataset = VariableLengthDataset(sequences, labels)
dataloader = DataLoader(
    dataset,
    batch_size=2,
    shuffle=True,
    collate_fn=custom_collate  # 关键配置
)

# 测试批次输出
for batch in dataloader:
    seqs, lbls, lens = batch
    print("Padded sequences:\n", seqs)
    print("Labels:", lbls)
    print("Lengths:", lens)

最佳实践

• 填充值选择：使用不会出现在实际数据中的值（如0）
• 排序重要性：提升RNN计算效率（避免过多填充计算）
• 长度张量：LSTM需传入lengths作为enforce_sorted=False参数
• 内存优化：大数据集使用__getitem__中即时加载数据

常见错误

• 忘记排序：导致RNN计算效率低下
• 错误填充方向：左侧/右侧填充需与模型设计一致
• 长度信息丢失：未将长度信息传递给模型
• 张量类型不匹配：确保dtype与模型权重一致

扩展知识

• Pack/Pad原理：PyTorch的pack_padded_sequence和pad_packed_sequence函数可跳过填充位置计算
• 掩码处理：在Attention机制中需使用attention_mask忽略填充位置
• 性能优化：使用pin_memory=True加速GPU数据传输
• 替代方案：Hugging Face的DataCollatorForTokenClassification等预置工具