pytorch版Unet的训练问题一

原创

于 2021-01-14 10:35:41 发布 · 3k 阅读

CC 4.0 BY-SA版权

博客作者在使用Pytorch实现UNet模型时遇到‘only batches of spatial targets supported’错误。问题源于目标张量维度不匹配，官方代码在处理输入和输出时存在不足。作者创建了自己的代码分支来改进，并找到了解决方案：针对网络类别数，调整训练和评估过程中的损失计算方式。具体改动包括在训练脚本中根据网络类别数决定是否挤压目标张量的维度，以及在评估脚本中处理单个样本的损失计算。

问题一错误内容 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"