pytorch版Unet的训练问题一

问题一 错误内容 only batches of spatial targets supported (3D tensors) but got targets of dimension: 4

pytorch版Unet 还是非常简洁易懂的.
官方给的代码有点小问题, 在输入和输出方面不是很好处理.
为了方便自己, 而且我打算日后改一下网络模型. 所以我自己在官方的示例上建了一个分支
官方的分支是
https://github.com/milesial/Pytorch-UNet
我的分支是
https://github.com/phker/Pytorch-MyUNet
目前代码还没提交.
训练的时候遇到了下面这个问题.

Traceback (most recent call last):
  File "f:/project/AI/123/AILabelSystem/Server/Pytorch-MyUNet/mytrain.py", line 216, in <module>
    train_net(net=net,
  File "f:/project/AI/123/AILabelSystem/Server/Pytorch-MyUNet/mytrain.py", line 89, in train_net
    loss = criterion(masks_pred, true_masks) # 求损失
  File "D:\Anaconda3\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "D:\Anaconda3\lib\site-packages\torch\nn\modules\loss.py", line 961, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "D:\Anaconda3\lib\site-packages\torch\nn\functional.py", line 2468, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "D:\Anaconda3\lib\site-packages\torch\nn\functional.py", line 2266, in nll_loss
    ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4

解决方案原文在这里.
https://github.com/milesial/Pytorch-UNet/issues/123

将 train.py

loss = criterion(masks_pred, true_masks)  # 80行上下

改成下面的这种

  if net.n_classes > 1: 
     loss = criterion(masks_pred, true_masks.squeeze(1)) # 求损失  # patch 123#bug
  else:
     loss = criterion(masks_pred, true_masks) # 求损失

同时修改eval.py将

if net.n_classes > 1:
  # tot += F.cross_entropy(mask_pred, true_masks).item() # patch 123#bug  把这行注释掉, 改成下面这行
    tot += F.cross_entropy(mask_pred.unsqueeze(dim=0), true_masks.unsqueeze(dim=0).squeeze(1)).item()  # patch 123#bug
else:
    pred = torch.sigmoid(mask_pred)
    pred = (pred > 0.5).float()
    tot += dice_coeff(pred, true_masks).item()