【解决办法】torch交叉熵使用时遇到 Dimension out of range

绝地灬酷狼 2022-01-30 15:27 830阅读 0赞

## 简述 ##

其实这个问题我很久以前用pytorch写程序的时候就遇到过这个问题，当时纠结了很久之后最后解决了。当时本来就想来写个东西来记录下避免其他人也遇到这样的问题。但后面我个菜鸡就完全忘记了emmmm

## 展示说明 ##

不同于我以往直接给出结论。这次我会通过展示的方式较为详细的讲解下。

*  导入包

import torch
    import torch.nn as nn

*  创造这个损失函数类

loss  = nn.CrossEntropyLoss()

*  随机生成下这个数据A

A = torch.randn(3, 5, requires_grad=True)

内容如下：

A
    tensor([[ 1.0483, -1.4251,  1.0502,  0.2437,  0.1477],
            [-0.1300, -1.0798, -1.0835, -1.0473,  1.2076],
            [-0.1647, -1.1366,  0.7088,  1.3719,  1.9135]], requires_grad=True)

*  随机生成3个label

target = torch.empty(3, dtype=torch.long).random_(5)

内容如下：

target
    tensor([2, 2, 3])

### 没有问题 ###

output = loss(A, target)

*  这时候就是没有问题的

output
    tensor(1.6934, grad_fn=<NllLossBackward>)

### 有问题的情况演示 ###

*  将前面创建的A在一维上取argmax。得到对应的label

A_arg = torch.argmax(A, dim=1)

内容如下：

A_arg
    tensor([2, 4, 4])

*  调用，出现题目所示的bug

output = loss(A_arg, target)

报错非常一大串emmm。关键就是下面这句。

RuntimeError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

## 解释 ##

其实就是因为torch的交叉熵的输入第一个位置的输入应该是在每个label下的概率, 而不是对应的label。

所以直接写成label的你，就出现上面所说的错误了。

NLLoss还是CE两者都是类似的。第一个输入为每个样本在不同label下的概率值（矩阵），第二个输入为label。注意，这里的label本质上都是概率矩阵在列方向上的index（也就是对应于某个样本的index类别）

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2ExOTk5MDQxMg_size_16_color_FFFFFF_t_70]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2ExOTk5MDQxMg_size_16_color_FFFFFF_t_70]: /images/20220130/bd045039395c4a7e9b2f65410d8ef458.png