PyTorch RuntimeError求助：张量维度不匹配问题修复

出错代码示例 #

import torch
import torch.nn as nn

class upd_GELU(nn.Module):
    def __init__(self):
        super().__init__()
        # 错误：参数形状设为[2]，但输入张量最后一维为120000
        self.alpha = nn.Parameter(torch.randn(2))
        self.beta = nn.Parameter(torch.randn(2))
    
    def forward(self, x):
        # 此处触发维度不匹配错误：x的维度2是120000，alpha/beta的维度0是2
        return x * torch.sigmoid(self.alpha * x + self.beta)

# 模拟输入：形状为[batch=32, seq_len=10, feat_dim=120000]
x = torch.randn(32, 10, 120000)
gelu = upd_GELU()
output = gelu(x)

完整报错堆栈示例 #

RuntimeError: The size of tensor a (120000) must match the size of tensor b (2) at non-singleton dimension 2
Stack trace:
...
File "your_script.py", line X, in forward
    return x * torch.sigmoid(self.alpha * x + self.beta)

修复方法 #

方案1：让参数形状匹配输入最后一维 #

如果需要对每个特征维度单独学习参数，初始化时将参数形状设为输入张量的最后一维大小：

class upd_GELU(nn.Module):
    def __init__(self, feat_dim=120000):
        super().__init__()
        self.alpha = nn.Parameter(torch.randn(feat_dim))
        self.beta = nn.Parameter(torch.randn(feat_dim))
    
    def forward(self, x):
        # alpha/beta形状[120000]，与x的最后一维匹配，可自动广播
        return x * torch.sigmoid(self.alpha * x + self.beta)

方案2：使用单维度参数实现全局共享 #

如果不需要按特征维度单独学习参数，将参数初始化为单维度（大小1），触发广播机制：

class upd_GELU(nn.Module):
    def __init__(self):
        super().__init__()
        self.alpha = nn.Parameter(torch.randn(1))
        self.beta = nn.Parameter(torch.randn(1))
    
    def forward(self, x):
        # alpha/beta形状[1]，可广播到x的所有维度
        return x * torch.sigmoid(self.alpha * x + self.beta)

方案3：手动扩展参数维度（适用于特殊场景） #

如果参数形状固定为[2]但需要与输入匹配，手动扩展参数维度以满足广播条件：

class upd_GELU(nn.Module):
    def __init__(self):
        super().__init__()
        self.alpha = nn.Parameter(torch.randn(2))
        self.beta = nn.Parameter(torch.randn(2))
    
    def forward(self, x):
        # 将参数扩展为[1,2,1]，匹配x的[32,2,120000]维度
        alpha_expanded = self.alpha.unsqueeze(0).unsqueeze(-1)
        beta_expanded = self.beta.unsqueeze(0).unsqueeze(-1)
        return x * torch.sigmoid(alpha_expanded * x + beta_expanded)

为何未触发张量广播机制 #

PyTorch广播机制的核心规则是：

1. 从最后一维开始向前匹配张量维度
2. 每个维度的大小必须满足：要么完全相等，要么其中一个张量的该维度大小为1
3. 若张量维度数量不同，会自动在前面补1维直到维度数量一致

你的报错场景中，输入张量x的维度2大小是120000，参数alpha/beta的对应维度大小是2——两个维度大小都大于1且不相等，完全不符合广播规则，因此无法触发广播，直接抛出维度不匹配错误。

内容的提问来源于stack exchange，提问作者Da-qiong