如何在PyTorch中针对Fashion-MNIST正确使用MSELoss并获取网络前两层的均方误差梯度？

首先，咱们先拆解你遇到的报错原因，再一步步调整代码满足你的需求。

错误根源：张量尺寸不匹配 #

你看到的RuntimeError: The size of tensor a (128) must match the size of tensor b (96)，本质是这两个问题：

1. model[0].forward(images)输出的是整个批次的特征，形状是(batch_size, 128)（每个样本对应128维第一层输出），但你取了output[0]，只拿到第一个样本的128维向量。
2. 你的labels是当前批次的类别标签，形状是(batch_size,)，而Fashion-MNIST训练集最后一个批次的样本数是96（60000 ÷ 128 = 468余96），所以96长度的标签和128长度的单样本特征，维度完全不匹配，导致MSELoss计算失败。

另外，你的需求是查看第一层与第二层之间的输出梯度，但直接用类别标签和128维中间特征计算MSE是不合理的——类别标签是0-9的离散值，和中间层的连续特征语义、维度都不匹配，我们需要调整目标张量的形状和含义。

修改方案：针对中间层计算MSE并查看梯度 #

假设你的需求是：计算**第一层Linear(784,128)的输出（也就是第二层ReLU的输入）**与目标张量的MSE，然后查看该中间层输出的梯度。以下是调整后的完整训练循环代码：

for e in range(epochs):
    running_loss = 0
    for images, labels in trainloader:
        # Flatten Fashion-MNIST images into a 784 long vector
        images = images.to(device)
        labels = labels.to(device)
        images = images.view(images.shape[0], -1)
        
        optimizer.zero_grad()
        # 获取第一层的输出：形状为(batch_size, 128)
        first_layer_output = model[0](images)
        # 构造和第一层输出形状完全匹配的目标张量，这里用全零张量做示例
        # 你也可以换成其他自定义目标，比如随机初始化的张量
        target = torch.zeros_like(first_layer_output).to(device)
        
        # 计算MSE损失：输入和目标形状完全一致，不会报错
        loss = criterion(first_layer_output, target)
        loss.backward()
        # 查看第一层输出的梯度（可选，按需打印）
        # print("第一层输出的梯度形状：", first_layer_output.grad.shape)
        
        optimizer.step()
        running_loss += loss.item()
    else:
        print(f"Epoch {e+1} 训练损失: {running_loss/len(trainloader):.3f}")
        # 保留原有的验证逻辑，但注意验证时要跑完整模型
        test_loss = 0
        accuracy = 0
        
        with torch.no_grad():
            model.eval()
            for images, labels in testloader:
                images = images.to(device)
                labels = labels.to(device)
                images = images.view(images.shape[0], -1)
                # 验证时跑完整模型得到分类输出
                ps = model(images)
                # 分类任务用NLLLoss更合理（配合模型最后的LogSoftmax）
                test_loss += F.nll_loss(ps, labels).item()
                
                top_p, top_class = ps.topk(1, dim=1)
                equals = top_class == labels.view(*top_class.shape)
                accuracy += torch.mean(equals.type(torch.FloatTensor)).item()
        
        model.train()
        print("Epoch: {}/{}..".format(e+1, epochs),
              "Training loss: {:.3f}..".format(running_loss/len(trainloader)),
              "Test loss: {:.3f}..".format(test_loss/len(testloader)),
              "Test Accuracy: {:.3f}".format(accuracy/len(testloader)))

额外说明 #

1. 梯度查看：在loss.backward()之后，你可以通过first_layer_output.grad直接访问第一层输出的梯度张量，查看其形状和数值。
2. 如果是分类任务想用MSELoss：如果你其实是想做分类但误用了中间层，那需要把类别标签转换成one-hot编码（形状(batch_size,10)），和模型最后一层的10维输出匹配，再计算MSE。
3. 损失选择：Fashion-MNIST是分类任务，优先用NLLLoss（配合LogSoftmax）或CrossEntropyLoss，MSELoss更适合回归任务。

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