nn.CrossEntropyLoss 输入为[n,class_num,H,W]

import torch
from torch import nn
data=torch.arange(54,dtype=torch.float32).reshape(2,3,3,3)
print(data)
x_l=torch.tensor([[[1,2,0],
                [1,2,0],
                [0,0,0]],

                [[1,2,0],
                [1,2,0],
                [0,0,0]]])
#[2,3,3,2]->[2,3,2,3] [b.类别，行，列]->[b,行，列，3]
data_permute=torch.permute(data,(0,2,3,1))
# data_permute=data
# print(data_permute)
data_permute_shape=data_permute.reshape(-1,3)
print(data_permute_shape)
#[b,3,2]->[b,-1]
x_l_shape=x_l.reshape(-1)
loss_fn = nn.CrossEntropyLoss()
loss_2D=loss_fn(data_permute_shape,x_l_shape)
print(loss_2D)
loss_3D=loss_fn(data,x_l)
print(loss_3D)
#[n，class_num] #[n]
#[n,class,H,W] #[n,H,W]