Is the calculation of mean IoU reasonable?

❓ Questions & Help

in the function of mean_iou, I mean the code “iou[torch.isnan(iou)] = 1”, is it reasonable,? when I test my model, some type has a bad segmentation, like board on s3dis dataset, it should get a very low score, but get a high score(0.95), in IOU, there are many ‘nan’ value，the code change them to 1, so it’s very high score. so is it reasonable to change ‘nan’ to 1?

def mean_iou(pred, target, num_classes, batch=None): r""“Computes the mean intersection over union score of predictions. Args: pred (LongTensor): The predictions. target (LongTensor): The targets. num_classes (int): The number of classes. batch (LongTensor): The assignment vector which maps each pred-target pair to an example. :rtype: :class:Tensor “”” i, u = intersection_and_union(pred, target, num_classes, batch) iou = i.to(torch.float) / u.to(torch.float) iou[torch.isnan(iou)] = 1 //in multiple class segmentation task, some types are not include in the test data iou = iou.mean(dim=-1) return iou

[GericoVi]

Yes I could do that. If the below is acceptable, I will do a pull request.

def mean_iou(pred, target, num_classes, batch=None, omitnans=False):
    r"""Computes the mean intersection over union score of predictions.

    Args:
        pred (LongTensor): The predictions.
        target (LongTensor): The targets.
        num_classes (int): The number of classes.
        batch (LongTensor): The assignment vector which maps each pred-target
            pair to an example.

    :rtype: :class:`Tensor`
    """
    i, u = intersection_and_union(pred, target, num_classes, batch)
    iou = i.to(torch.float) / u.to(torch.float)

    if omitnans:
        iou = iou[~iou.isnan()].mean()
    else:
        iou[torch.isnan(iou)] = 1
        iou = iou.mean(dim=-1)
    
    return iou

For reference, on a trained PointNet++ segmentation model, an example from the airplane class of ShapeNet gives an mIoU of 0.85191 with omitnans = True and 0.98815 with omitnans = False. The former agrees with published results better.

[rusty1s]

How about we add a flag to let the user decide on how to incorporate NaN values? How are other libraries handling thise use-case? Adding support for torch.nanmean() sounds great IMO.