多头注意力模型训练报错：IndexError维度越界问题求助

错误根源 #

你碰到的IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)本质是张量维度不匹配，结合给出的张量信息：

• y_pred是一维张量（shape [3]），但CrossEntropyLoss要求模型输出必须是(batch_size, num_classes)的二维张量
• y是0维标量，而损失函数期望目标标签是(batch_size)的一维张量
• 更深层问题是多头注意力层的维度处理错误，且表格数据的输入维度没有适配注意力层的要求（注意力需要带序列维度的张量）

具体问题点拆解 #

1. 输入维度不兼容：你的表格数据样本是(input_dim)的一维张量，批量后是(batch_size, input_dim)，但多头注意力需要(batch_size, seq_len, input_dim)格式的输入（需包含序列维度）。
2. 注意力层维度计算错误：MultiheadAttention的forward方法中，对query/key/value的reshape逻辑错误，没有正确拆分多头维度，且缩放因子用了input_dim而非单头维度，导致后续张量维度混乱。
3. 过度squeeze操作：训练循环中y_pred = model(x).squeeze()把本该是(batch_size, 3)的输出挤成一维，彻底破坏了损失函数要求的维度匹配。

分步修复 #

1. 给输入增加序列维度 #

在geneGPT的forward方法中，先给输入补上序列维度（把每个样本看作长度为1的序列），处理完注意力后再移除：

class geneGPT(nn.Module):
    def __init__(self, input_dim, hid_dim, output_dim, num_heads, dropout_rate):
        super().__init__()
        self.attention = MultiheadAttention(input_dim, num_heads, dropout_rate)
        self.fc1 = nn.Linear(input_dim, hid_dim)  # 修正输入维度，和注意力输出对齐
        self.relu = nn.ReLU()
        self.out = nn.Linear(hid_dim, output_dim)

    def forward(self, x, mask=None):
        x = x.unsqueeze(1)  # 从(batch_size, input_dim) → (batch_size, 1, input_dim)
        x = self.attention(x, x, x, mask)
        x = x.squeeze(1)  # 从(batch_size, 1, input_dim) → (batch_size, input_dim)
        x = self.relu(self.fc1(x))
        x = self.out(x)
        return x

2. 修复多头注意力层的维度逻辑 #

修正MultiheadAttention的维度拆分、注意力计算和拼接逻辑：

class MultiheadAttention(nn.Module):
    def __init__(self, input_dim, num_heads, dropout_rate):
        super(MultiheadAttention, self).__init__()
        self.input_dim = input_dim
        self.num_heads = num_heads
        self.head_dim = input_dim // num_heads  # 单头维度，必须能被input_dim整除
        self.dropout_rate = dropout_rate

        self.q_linear = nn.Linear(input_dim, input_dim)
        self.k_linear = nn.Linear(input_dim, input_dim)
        self.v_linear = nn.Linear(input_dim, input_dim)
        self.dropout = nn.Dropout(dropout_rate)
        self.out_linear = nn.Linear(input_dim, input_dim)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        seq_len = query.size(1)

        # 线性变换后拆分多头：(batch_size, seq_len, input_dim) → (batch_size, num_heads, seq_len, head_dim)
        q = self.q_linear(query).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_linear(key).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_linear(value).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算缩放点积注意力，用单头维度做缩放因子
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        
        if mask is not None:
            mask = mask.unsqueeze(1).unsqueeze(2)
            scores = scores.masked_fill(mask == 0, -1e9)

        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 计算注意力输出并拼接多头
        attn_output = torch.matmul(attn_weights, v)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.input_dim)
        
        out = self.out_linear(attn_output)
        return out

3. 移除过度的squeeze操作 #

训练循环中保留模型输出的正确维度，仅在必要时处理标签：

# 训练循环内的修改
y_pred = model(x)  # 输出形状为(batch_size, 3)
# 仅当y是(batch_size,1)时才需要squeeze，否则直接使用
if y.dim() > 1:
    y = y.squeeze()

修复后验证 #

• 输入x：(batch_size, input_dim) → 补序列维度后(batch_size, 1, input_dim)
• 注意力输出：(batch_size, 1, input_dim) → 移除序列维度后(batch_size, input_dim)
• 最终模型输出：(batch_size, 3)，标签y为(batch_size)，完全符合CrossEntropyLoss的维度要求

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