PyTorch中可学习加权均值操作的实现问题求助

错误原因 #

你之前的代码报错是因为nn.Linear(max_length, embedding_dim)的使用逻辑错误：Linear层默认对输入的最后一个维度做线性变换，但你的输入在embedding后是(batch_size, max_length, embedding_dim)，此时Linear会错误地对每个embedding_dim维度的元素做变换，导致维度不匹配，最终触发形状相乘错误。

正确实现思路 #

要实现可学习的加权均值，核心是定义一个形状为(max_length,)的可训练参数，将其与embedding后的张量做加权求和，替代等权的torch.mean。为了让权重更合理，通常会对权重做softmax归一化，确保权重非负且总和为1（也可根据需求跳过归一化）。

完整代码实现 #

import torch
import torch.nn as nn

class EmbeddingClassifier(nn.Module):
    def __init__(self, vocab_size, embedding_dim, num_classes, max_length=20):
        super().__init__()
        self.max_length = max_length
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        
        # 定义可学习的权重参数，形状为(max_length,)
        self.weights = nn.Parameter(torch.randn(max_length))
        
        self.fc = nn.Linear(embedding_dim, num_classes)

    def forward(self, x):
        """
        x: pytorch LongTensor of input_ids, shape (batch_size, sequence_length)
        """
        # 补零到max_length，确保和输入x的设备一致（CPU/GPU）
        batch_size, seq_len = x.shape
        if seq_len < self.max_length:
            pad_len = self.max_length - seq_len
            padding = torch.zeros((batch_size, pad_len), dtype=x.dtype, device=x.device)
            x = torch.cat([x, padding], dim=1)
        
        # 经过embedding后，形状为(batch_size, max_length, embedding_dim)
        x_emb = self.embedding(x)
        
        # 对权重做softmax归一化，保证权重非负且总和为1
        normalized_weights = torch.softmax(self.weights, dim=0)
        # 扩展权重形状为(1, max_length, 1)，利用广播机制和embedding张量做逐元素乘法
        normalized_weights = normalized_weights.unsqueeze(0).unsqueeze(-1)
        
        # 在序列维度(dim=1)上加权求和，得到(batch_size, embedding_dim)的张量
        weighted_avg = torch.sum(x_emb * normalized_weights, dim=1)
        
        # 全连接层输出分类结果
        out = self.fc(weighted_avg)
        return out

关键细节说明 #

• 设备对齐：补零操作时使用x.device，确保padding张量和输入在同一设备，避免跨设备运算错误。
• 权重归一化：torch.softmax处理后的权重更贴近普通均值的逻辑，训练时更稳定；若不需要约束，可直接使用原始权重，但可能出现权重波动过大的情况。
• 维度广播：将(max_length,)的权重扩展为(1, max_length, 1)，无需手动循环即可实现批量样本的加权求和，符合PyTorch的张量运算逻辑。

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