具有交叉熵损失函数的 UNET

问题描述

我试图训练 UNET,输入大小为 [3,128,128],相应的掩码是 [1,128],它直接包含类(而不是像素,它将包含类号 - 1,2)。我正在尝试解决一个两类问题,因此我的掩码包含 1,2 作为标签。现在我将我的图像发送到模型,预测掩码的维度是 [2,128]。现在要训练模型,我选择 16 作为批量大小。所以,现在我输入为 [16,3,128] 所以预测维度是 [16,2,128]。但是我有真实的掩码 [16,1,128]。现在如何在 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,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,128])
shape of groundtruth masks: torch.Size([16,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,_Reduction.get_enum(reduction),ignore_index)
   2386     elif dim == 4:
-> 2387         ret = torch._C._nn.nll_loss2d(input,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,128]

您能否提出错误在哪里以及这个 CrossEntropyLoss() 在 Pytorch 中如何工作以进行图像分割?我在这里缺少什么?我曾尝试将目标大小重塑为 [16,128],这导致了另一个错误。非常感谢!

解决方法

documentation 指定如果输入是形状 (N,C,d1,d2),那么目标必须是形状 (N,d2)。相反,您的目标是形状 (N,1,d2),因此您需要删除不必要的单一维度。

loss = criterion(output,labels.squeeze(1))

如果您从此更改中收到另一个错误,则说明您的代码存在另一个问题,但这是 CrossEntropyLoss 的正确张量形状。