如何让PyTorch DataLoader对面板数据按个体而非观测值采样？

针对面板数据（单个体对应多条观测）在DataLoader中按个体采样的需求，以下是完整的修改方案，解决默认collate_fn无法处理变长张量的问题：

1. 重构ChoiceDataset，按个体索引返回数据 #

首先要让Dataset以个体为单位进行索引，而非单条观测：

• 提前整理individual_map：一个列表，每个元素对应一个个体的所有观测在原始数据中的行索引。
• 修改__len__返回个体总数，__getitem__根据个体索引提取该个体的全部观测数据。

示例代码：

import torch
from torch.utils.data import Dataset

class ChoiceDataset(Dataset):
    def __init__(self, raw_features, raw_labels, individual_map):
        # raw_features: 原始所有观测的特征张量，形状[N, 特征维度D]
        # raw_labels: 原始所有观测的标签张量，形状[N]
        # individual_map: 列表，每个元素是对应个体的观测索引列表
        self.raw_features = raw_features
        self.raw_labels = raw_labels
        self.individual_map = individual_map

    def __len__(self):
        # 返回个体总数，而非观测总数
        return len(self.individual_map)

    def __getitem__(self, idx):
        # 获取第idx个个体的所有观测索引
        obs_idx = self.individual_map[idx]
        # 提取该个体的全部特征和标签
        indiv_feat = self.raw_features[obs_idx]  # 形状[T, D]，T是该个体的观测数
        indiv_label = self.raw_labels[obs_idx]    # 形状[T]
        return indiv_feat, indiv_label

2. 自定义collate_fn处理变长张量 #

默认collate_fn会强制拼接不同长度的张量导致报错，这里提供两种实用的处理方式：

方式一：Padding到批次最大长度（适合固定输入形状的模型） #

将批次内所有个体的特征/标签padding到当前批次的最大观测数，同时记录每个个体的实际长度，方便后续过滤padding部分。

def custom_collate_pad(batch):
    # batch是列表，每个元素为(个体特征张量, 个体标签张量)
    feat_list, label_list = zip(*batch)
    
    # 统计批次内每个个体的观测长度
    lengths = torch.tensor([f.shape[0] for f in feat_list], dtype=torch.int64)
    max_seq_len = lengths.max().item()
    feat_dim = feat_list[0].shape[1]
    
    # 初始化padding后的张量
    padded_feats = torch.zeros((len(batch), max_seq_len, feat_dim), dtype=feat_list[0].dtype)
    padded_labels = torch.zeros((len(batch), max_seq_len), dtype=label_list[0].dtype)
    
    # 填充数据
    for i, (feat, label) in enumerate(zip(feat_list, label_list)):
        padded_feats[i, :len(feat)] = feat
        padded_labels[i, :len(label)] = label
    
    return padded_feats, padded_labels, lengths

方式二：保留变长张量列表（适合支持变长输入的模型） #

如果模型支持变长输入（比如LSTM配合pack_padded_sequence），可以直接返回变长张量列表，同时附带长度信息：

def custom_collate_list(batch):
    feat_list, label_list = zip(*batch)
    lengths = torch.tensor([f.shape[0] for f in feat_list], dtype=torch.int64)
    # 按长度降序排序（pack_padded_sequence要求输入从长到短排列）
    sorted_idx = torch.argsort(lengths, descending=True)
    sorted_feats = [feat_list[i] for i in sorted_idx]
    sorted_labels = [label_list[i] for i in sorted_idx]
    sorted_lengths = lengths[sorted_idx]
    return sorted_feats, sorted_labels, sorted_lengths

3. 初始化DataLoader时指定自定义collate_fn #

把自定义的collate_fn传入DataLoader，此时batch_size指的是每个批次的个体数量：

# 假设已准备好raw_features, raw_labels, individual_map
dataset = ChoiceDataset(raw_features, raw_labels, individual_map)
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    collate_fn=custom_collate_pad  # 或custom_collate_list，根据模型需求选择
)

4. 模型内的适配处理 #

• 如果用了padding方案：计算损失时需要过滤padding部分，比如假设padding值为0，用mask屏蔽无效标签：

  # 模型输出形状：[batch_size, max_seq_len, num_classes]
  outputs = model(padded_feats, lengths)
  # 生成mask，过滤padding的标签
  mask = (padded_labels != 0)
  loss = criterion(outputs[mask], padded_labels[mask])
  

• 如果用了变长列表方案：在模型中用pack_padded_sequence处理输入：

  sorted_feats, sorted_labels, sorted_lengths = batch
  # 打包变长序列
  packed_feats = torch.nn.utils.rnn.pack_sequence(sorted_feats)
  packed_output, _ = lstm(packed_feats)
  # 解包为padding后的张量（可选）
  padded_output, _ = torch.nn.utils.rnn.pad_packed_sequence(packed_output)
  

内容的提问来源于stack exchange，提问作者Álvaro A. Gutiérrez-Vargas