pytorch学习笔记（三）用nn.cross_entropy所遇到的问题

最新推荐文章于 2025-02-18 20:29:03 发布

原创

最新推荐文章于 2025-02-18 20:29:03 发布 · 4.5k 阅读

0 ·

CC 4.0 BY-SA版权

本文通过Mnist数据集实例，探讨使用PyTorch中nn.CrossEntropyLoss时遇到的问题。在训练过程中，由于实际输出yb与模型预测输出model(xb)尺寸不匹配，导致RuntimeError。解决方法是确保yb和model(xb)尺寸一致。

Mnist数据集的图片尺寸一般是28*28 = 784,这里模拟一下搭建一个神经元的网络来进行预测和loss的计算，代码如下。

import torch
from torch import nn
import math
import torch.nn.functional as F

class Mnist_net(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(784,10)/math.sqrt(784))
        self.bias = nn.Parameter(torch.zeros(10))

    def forward(self,xb):
        xb = xb@self.weights +self.bias
        return xb

model = Mnist_net()
loss_func = F.cross_entropy
xb = torch.randn(2,784)
yb = torch.tensor([1,1])
print ("model(xb):",model(xb))
print ("yb:",yb)
loss = loss_func(model(xb),yb)
print ("loss:",loss)

这里要注意的是yb是实际输出的标注，model（xb）是预测的输出，这里yb和model（xb）是不同的尺寸，如下。

model(xb): tensor([[ 0.3825, -0.0071,  0.4110, -0.0534, -0.3756,  0.7349, -3.1753,  0.7359,
         -0.1801,  0.1951],
        [-1.4873,  0.8943, -0.4351, -0.0579,  0.3246,  1.3478, -1.2897,  1.0348,
         -0.8285,  0.8874]], grad_fn=<AddBackward0>)
y