如何调整collate_fn实现固定长度填充，解决嵌入维度动态变化问题

问题背景 #

我想要训练一个以embedding_dim为参数的简单神经网络，模型定义如下：

class BoolQNN(nn.Module):
    def __init__(self, embedding_dim):
        super(BoolQNN, self).__init__()
        self.fc1 = nn.Linear(embedding_dim, 64)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(64, 1)

    def forward(self, question_emb, passage_emb):
        combined = torch.cat((question_emb, passage_emb), dim=1)
        x = self.fc1(combined)
        x = self.relu(x)
        x = self.fc2(x)
        return torch.sigmoid(x)

使用PyTorch的DataLoader结合自定义collate_fn加载数据：

train_dataset = BoolQDataset(train_data, pretrained_embeddings)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True,collate_fn=collate_fn_padd)

model = BoolQNN(301)

collate_fn_padd函数实现如下：

def collate_fn_padd(batch):

  questions, passages, labels = zip(*batch)

  questions = [torch.tensor(q) for q in questions]
  passages = [torch.tensor(p) for p in passages]

  padded_questions = pad_sequence(questions, batch_first=True, padding_value=0)
  padded_passages = pad_sequence(passages, batch_first=True, padding_value=0)

  labels = torch.tensor(labels, dtype=torch.float32)
  
  return padded_questions, padded_passages, labels

问题：当前每个训练批次会将文本嵌入填充至该批次的最长序列长度，导致神经网络线性层的输入维度（嵌入维度）随批次动态变化，与模型初始化时固定的embedding_dim不匹配，进而出现如“mat1 and mat2 shapes cannot be multiplied (16x182 and 301x64)”的报错。需要调整collate_fn_padd函数，实现与批次无关的固定长度填充。

解决方案 #

要实现固定长度填充，需提前确定全局固定的最大序列长度（可根据全数据集统计结果设定，或直接指定符合需求的固定值），然后在collate_fn中对每个样本进行截断或填充，统一到该固定长度。

修改后的collate_fn_padd代码 #

def collate_fn_padd(batch, fixed_q_len=150, fixed_p_len=151):
    # 注意：fixed_q_len + fixed_p_len 必须等于模型初始化的embedding_dim(301)
    questions, passages, labels = zip(*batch)

    processed_questions = []
    for q in questions:
        q_tensor = torch.tensor(q)
        # 截断过长序列，填充过短序列
        if len(q_tensor) > fixed_q_len:
            q_tensor = q_tensor[:fixed_q_len]
        else:
            pad_len = fixed_q_len - len(q_tensor)
            q_tensor = torch.cat([q_tensor, torch.zeros(pad_len, dtype=q_tensor.dtype)])
        processed_questions.append(q_tensor)
    
    processed_passages = []
    for p in passages:
        p_tensor = torch.tensor(p)
        if len(p_tensor) > fixed_p_len:
            p_tensor = p_tensor[:fixed_p_len]
        else:
            pad_len = fixed_p_len - len(p_tensor)
            p_tensor = torch.cat([p_tensor, torch.zeros(pad_len, dtype=p_tensor.dtype)])
        processed_passages.append(p_tensor)
    
    # 转换为batch张量
    padded_questions = torch.stack(processed_questions)
    padded_passages = torch.stack(processed_passages)
    labels = torch.tensor(labels, dtype=torch.float32)
    
    return padded_questions, padded_passages, labels

关键说明 #

1. 维度匹配：模型初始化时embedding_dim=301，对应问题和段落嵌入拼接后的总长度，因此必须保证fixed_q_len + fixed_p_len = 301，确保拼接后的维度固定，与fc1输入维度匹配。
2. 截断与填充逻辑：对每个样本，长度超过固定值则截断到固定长度，不足则用0填充（若使用预训练词嵌入，建议用预训练模型自带的padding token替代0，效果更优）。
3. DataLoader调用调整：使用时需传递固定长度参数，或直接在函数内写死：

train_loader = torch.utils.data.DataLoader(
    train_dataset, 
    batch_size=16, 
    shuffle=True,
    collate_fn=lambda batch: collate_fn_padd(batch, fixed_q_len=150, fixed_p_len=151)
)

额外注意点 #

• 若不确定固定长度，可先统计全数据集内问题和段落的最大长度，取该值作为固定长度，或在最大长度基础上增加少量冗余。
• 填充值可根据实际需求调整，比如预训练embedding的padding token对应的向量，而非简单的0向量。

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