Python 具有交叉熵损失函数的UNET

我试图训练UNET，输入大小为[3128128]，对应的掩码为[1128128]，它直接包含类（而不是像素，它将包含类号-1,2）。我试图解决一个两类问题，因此我的掩码包含1,2作为标签。现在我将图像发送到模型，预测的遮罩尺寸为[2128128]。现在要训练一个模型，我选择16作为批量大小。现在输入为[16,3128128]，所以预测的维数是[16,2128128]。但我有地面真相面具[16,1128128]。现在我如何在Pytorch中应用交叉熵损失？我尝试了以下方法并得到了以下错误。你能帮忙吗？提前谢谢

lr = 0.1  # 0.1
criterion = nn.CrossEntropyLoss() #nn.L1Loss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, nesterov=True, weight_decay=0.0001)
is_train = True
is_pretrain = False
acc_best = 0
total_epoch = 30
if is_train is True:
    # Training
    for epoch in range(total_epoch):
        model.train()
        tims = time.time()
        for i, (images, labels) in enumerate(train_loader):
            images = Variable(images.permute(0,3,1,2).float().cuda())                                                                                                   
            labels = Variable(labels.type(torch.LongTensor).cuda())
            # Forward + Backward + Optimize
            optimizer.zero_grad()
            outputs = model(images)
            outputs = outputs.type(torch.float)
            print('predictedLabelsType:',outputs[0].type())
            print('ActualLabelsType:',labels[0].type())
            print('shape of predicted outputs:',outputs.shape)
            print('shape of groundtruth masks:',labels.shape)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

我的输出如下：

predictedLabelsType: torch.cuda.FloatTensor
ActualLabelsType: torch.cuda.LongTensor
shape of predicted outputs: torch.Size([16, 2, 128, 128])
shape of groundtruth masks: torch.Size([16, 1, 128, 128])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-8-b692a8d536a9> in <module>()
     52             print('shape of predicted outputs:',outputs.shape)
     53             print('shape of groundtruth masks:',labels.shape)
---> 54             loss = criterion(outputs, labels)
     55             loss.backward()
     56             optimizer.step()

3 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2385         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2386     elif dim == 4:
-> 2387         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2388     else:
   2389         # dim == 3 or dim > 4

RuntimeError: 1only batches of spatial targets supported (3D tensors) but got targets of size: : [16, 1, 128, 128]

predictedLabelsType:torch.cuda.FloatTensor
实际标签类型：torch.cuda.LongTensor
预测输出的形状：火炬大小（[16,2,128,128]）
地面真相面具的形状：火炬。大小（[16,1,128,128]）
---------------------------------------------------------------------------
运行时错误回溯（上次最近调用）
在（）
52打印（'预测输出的形状：'，outputs.shape）
53打印（'groundtruth遮罩的形状：'，标签.形状）
--->54损耗=标准（输出、标签）
55.损失向后（）
56优化器。步骤（）
3帧
/nll_损耗中的usr/local/lib/python3.7/dist-packages/torch/nn/functional.py（输入、目标、重量、尺寸平均值、忽略索引、减少、减少）
2385 ret=火炬。火炬。nll损失（输入，目标，重量，减少。获取枚举（减少），忽略索引）
2386 elif dim==4：
->2387 ret=torch.\u C.\u nn.nll\u loss2d（输入、目标、权重、减少、获取枚举（减少）、忽略索引）
2388其他：
2389#dim==3或dim>4
RuntimeError:1只支持成批的空间目标（3D张量），但得到的目标大小为：[16,1,128,128]

您能告诉我错误在哪里，以及这个CrossEntropyLoss（）如何在Pytorch中用于图像分割吗？我错过了什么？我曾尝试将目标大小重塑为[16128128]，这会导致另一个错误。非常感谢

指定如果输入是shape

（N，C，d1，d2）

，则目标必须是shape

（N，d1，d2）

。相反，您的目标是形状

（N，1，d1，d2）

，因此您需要删除不必要的单一维度

loss=标准（输出、标签、挤压（1））

如果此更改导致另一个错误，那么代码中还有另一个问题，但这是

CrossEntropyLoss

的正确张量形状