目录
在运行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
参考自: