PyTorch线性模型矩阵乘法错误求助：维度不匹配问题

错误原因分析 #

你遇到的(32, 784)和(1, 10)张量维度不匹配问题，本质是线性层的输入特征数与前一层输出特征数不兼容：

• 输入经nn.Flatten()后形状为(32, 784)（32是batch size，784是28×28像素展开后的特征数）
• 第一个nn.Linear层的in_features参数被错误设置为1，导致该层权重矩阵形状为(1, hidden_units)，当hidden_units=10时，就会触发(32,784)与(1,10)的矩阵乘法维度不匹配错误。

具体修复步骤 #

1. 正确初始化模型参数
   KMNIST输入是28×28单通道图像，展开后共784个特征，初始化模型时需严格对应参数：

   • input传入784（输入特征总数）
   • hidden_units传入合理正整数（如128、256，无需乘以784）
   • output传入10（对应10个分类）

   示例代码：

   model = KMNIST_V1(input=784, hidden_units=128, output=10)
   

2. 验证输入数据维度
   确保输入张量形状符合预期：

   • 原始图像数据应为(batch_size, 1, 28, 28)（通道在前）或(batch_size, 28, 28, 1)（通道在后），nn.Flatten()会自动将其展开为(batch_size, 784)
   • 若输入已为展开后的(batch_size, 784)，可直接兼容模型

3. 确认层间参数连贯性
   你的模型代码本身逻辑正确，每层nn.Linear的in_features已对应前一层out_features，只需保证初始化参数传递正确即可。

测试修复后的完整代码 #

import torch
import torch.nn as nn

class KMNIST_V1(nn.Module):
  def __init__(self,
               input: int,
               hidden_units: int,
               output: int):
    super().__init__()
    self.layer_stack = nn.Sequential(
        nn.Flatten(),
        nn.Linear(in_features = input,
                  out_features = hidden_units),
        nn.Linear(in_features = hidden_units,
                  out_features = hidden_units),
        nn.Linear(in_features = hidden_units,
                  out_features = output)
    )

  def forward(self, x):
    return self.layer_stack(x)

# 初始化正确参数的模型
model = KMNIST_V1(input=784, hidden_units=128, output=10)
# 生成模拟输入（batch_size=32，1通道28×28图像）
test_input = torch.randn(32, 1, 28, 28)
# 前向传播测试
output = model(test_input)
print(output.shape)  # 应输出torch.Size([32, 10])

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