RuntimeError: 1only batches of spatial targets supported (non-empty 3D tensors) but got targets of s

目录

在运行UNet3+进行多分类训练的时候报错

RuntimeError: 1only batches of spatial targets supported (non-empty 3D tensors) but got targets of size xxx

1. 修改train.py

2. 修改predict.py

3. 修改eval.py

然后运行下面语句就可以开始训练了：

【实验记录】U-Net（PyTorch）

---

在运行UNet3+进行多分类训练的时候报错

RuntimeError: 1only batches of spatial targets supported (non-empty 3D tensors) but got targets of size xxx

经过多次查阅资料加试错后，得到了如下的解决办法：

UNet3+源代码对于做多分类任务有点小 bug，这里稍微修改一下。（UNet3+代码出自githubGitHub - avBuffer/UNet3plus_pth: UNet3+/ UNet++/UNet, used in Deep Automatic Portrait Matting in Pytorth）

1. 修改train.py

# line 56
if net.n_classes > 1:
        criterion = nn.CrossEntropyLoss()  # 使用多分类损失
    else:
        criterion = nn.BCEWithLogitsLoss()

# line 79
loss = criterion(masks_pred, torch.squeeze(true_masks))   # 修改

# line 153
net = UNet(n_channels=3, n_classes=3)   # 修改类别数

其中使用torch.squeeze()函数的原因是因为在使用Crossentropyloss作为损失函数时，output=net(input)的output应该是[batchsize, n_class, height, weight]，而label则是[batchsize, height, weight]，label是单通道灰度图;BCELoss与CrossEntropyLoss都是用于分类问题。BCELoss是CrossEntropyLoss的一个特例，只用于二分类问题，而CrossEntropyLoss可以用于二分类，也可以用于多分类。

nn.CrossEntropyLoss()函数计算交叉熵损失举例：

使用：
# output是网络的输出，size=[batch_size, class]
#如网络的batch size为128，数据分为10类，则size=[128, 10]
 
# target是数据的真实标签，是标量，size=[batch_size]
#如网络的batch size为128，则size=[128]
 
criterion=nn.CrossEntropyLoss()
crossentropyloss_output=criterion(output,target)

nn.CrossEntropyLoss()函数参考自：RuntimeError: 1only batches of spatial targets supported (non-empty 3D tensors) but got targets of size - DuanYongchun - 博客园

2. 修改predict.py

os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
if unet_type == 'v2':
    net = UNet2Plus(n_channels=3, n_classes=1)
elif unet_type == 'v3':
# line93
    net = UNet3Plus(n_channels=3, n_classes=20)   # 修改类别数
    #net = UNet3Plus_DeepSup(n_channels=3, n_classes=1)
    #net = UNet3Plus_DeepSup_CGM(n_channels=3, n_classes=1)
else:
    net = UNet(n_channels=3, n_classes=1)

3. 修改eval.py

for true_mask, pred in zip(true_masks, mask_pred):
    pred = (pred > 0.5).float()
    if net.n_classes > 1:
# line26
    tot += F.cross_entropy(pred.unsqueeze(dim=0), true_mask).item()
    # tot += F.cross_entropy(pred.unsqueeze(dim=0), true_mask.unsqueeze(dim=0)).item()
    else:
        tot += dice_coeff(pred, true_mask.squeeze(dim=1)).item()

然后运行下面语句就可以开始训练了：

python train.py -g 0 -u v3 -e 400 -b 2 -l 0.1 -s 0.5 -v 15.0

参考自：

【实验记录】U-Net（PyTorch）

【实验记录】U-Net（PyTorch）_yaoyz105-CSDN博客