catalogue

An error is reported when running unet3 + for multi classification training

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

1. Modify train.py

2. Modify predict.py

3. Modify eval.py

Then run the following statement to start the training:

[experiment record] u-net (pytorch)

An error is reported when running unet3 + for multi classification training

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

After many times of reference, trial and error, the following solutions are obtained:

Unet3 + source code has a small bug for multi classification tasks, which is slightly modified here. (unet3 + code comes from githubgithub – avbuffer/unet3plus_pth: unet3 +/UNET + +/UNET, used in deep automatic portal matching in Python)

1. Modify 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)   

The reason for using the torch. Squeeze() function is that when crossentropy is used as the loss function, the output of output = net (input) should be [batchsize, n_class, height, weight], while the label is [batchsize, height, weight], and the label is a single channel gray map; Both bceloss and cross-entropy loss are used for classification problems. Bceloss is a special case of cross-entropy loss, which is only used for binary classification problems, and cross-entropy loss can be used for binary classification or multi-classification.

NN. Crossentropyloss() function to calculate cross entropy loss example:

Use.
# output is the output of the network, size=[batch_size, class]
# If the batch size of the network is 128 and the data is divided into 10 classes, then size=[128, 10]
 
# target is the real label of the data, which is scalar, size=[batch_size]
# If the batch size of the network is 128, then size=[128]
 
criterion=nn.CrossEntropyLoss()
crossentropyloss_output=criterion(output,target)

2. Modify 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. Modify 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()

Then run the following statement to start the training:

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