PyTorch张量维度不匹配错误求助：张量a(61)与张量b(64)维度3不匹配

核心问题 #

你遇到的RuntimeError: The size of tensor a (61) must match the size of tensor b (64) at non-singleton dimension 3，本质是模型输入输出维度不匹配，叠加标签处理逻辑错误、损失函数选型不当导致的连锁问题。

具体修复步骤 #

1. 修正模型输入层（适配图片维度） #

你的模型第一层是nn.Linear(28, 16)，但CHoiCe-Dataset的图片是28×28的二维灰度图，需要先展平为一维向量（28×28=784）才能输入全连接层。修改模型：

model = nn.Sequential(
    nn.Linear(28*28, 16),  # 替换原28为28*28，匹配展平后的图片维度
    nn.Sigmoid(),
    nn.Linear(16, 16),
    nn.Sigmoid(),
    nn.Linear(16, 61)  # 61对应数据集类别数，此部分无需修改
)

2. 修正训练循环中的输入与标签处理 #

• 输入图片必须展平后传入模型
• 分类任务不要用MSELoss（回归场景损失），改用CrossEntropyLoss（分类任务专用，无需手动转换one-hot标签）
• 标签不需要view(1, -1)变形，直接使用原始类别索引即可（需确保标签为LongTensor类型）

修改后的训练函数：

def train():
    criterion = nn.CrossEntropyLoss()  # 替换原MSELoss
    optimizer = optim.SGD(model.parameters(), lr=0.3, momentum=0.9)

    for epoch in range(3000):
        running_loss = 0
        for images, labels in dataloader:
            optimizer.zero_grad()
            
            # 展平图片：将(batch_size, 1, 28, 28)转为(batch_size, 784)
            images_flat = images.view(images.size(0), -1)
            output = model(images_flat)
            
            # 确保标签为LongTensor类型，适配CrossEntropyLoss要求
            labels = labels.long()
            loss = criterion(output, labels)

            loss.backward()
            optimizer.step()
            running_loss += loss.item()

3. 验证Dataset的输出格式 #

确保自定义Dataset返回的内容符合要求：

• images是形状为(1, 28, 28)的张量（单通道灰度图）
• labels是0~60之间的整数（对应61个类别，与模型最后一层输出维度匹配）

额外建议 #

• 初始学习率0.3过高，建议先设置为0.01，后续可通过学习率调度器动态调整
• 3000轮epoch易导致过拟合，建议加入验证集监控损失变化，触发提前停止训练逻辑

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