Binary Segmentation with Mask Cross Entropy

[qritive-saran]

I was trying to use mask_cross_entropy function for a binary segmentation problem.
But I was getting errors related to the tensor datatype. The output tensor was float32 , but the target tensor was integer.
From the implementation, we see that the binary_cross_entropy function has converted the label to float tensor ,but the mask_cross_entropy function has not.

Can someone throw more light on why there is this difference or is it a bug ?

https://github.com/open-mmlab/mmsegmentation/blob/e64548fda0221ad708f5da29dc907e51a644c345/mmseg/models/losses/cross_entropy_loss.py#L157

`def binary_cross_entropy(pred,
label,
weight=None,
reduction='mean',
avg_factor=None,
class_weight=None,
ignore_index=-100,
avg_non_ignore=False,
**kwargs):
"""Calculate the binary CrossEntropy loss.

Args:
    pred (torch.Tensor): The prediction with shape (N, 1).
    label (torch.Tensor): The learning label of the prediction.
        Note: In bce loss, label < 0 is invalid.
    weight (torch.Tensor, optional): Sample-wise loss weight.
    reduction (str, optional): The method used to reduce the loss.
        Options are "none", "mean" and "sum".
    avg_factor (int, optional): Average factor that is used to average
        the loss. Defaults to None.
    class_weight (list[float], optional): The weight for each class.
    ignore_index (int): The label index to be ignored. Default: -100.
    avg_non_ignore (bool): The flag decides to whether the loss is
        only averaged over non-ignored targets. Default: False.
        `New in version 0.23.0.`

Returns:
    torch.Tensor: The calculated loss
"""
if pred.size(1) == 1:
    # For binary class segmentation, the shape of pred is
    # [N, 1, H, W] and that of label is [N, H, W].
    # As the ignore_index often set as 255, so the
    # binary class label check should mask out
    # ignore_index
    assert label[label != ignore_index].max() <= 1, \
        'For pred with shape [N, 1, H, W], its label must have at ' \
        'most 2 classes'
    pred = pred.squeeze()
if pred.dim() != label.dim():
    assert (pred.dim() == 2 and label.dim() == 1) or (
            pred.dim() == 4 and label.dim() == 3), \
        'Only pred shape [N, C], label shape [N] or pred shape [N, C, ' \
        'H, W], label shape [N, H, W] are supported'
    # `weight` returned from `_expand_onehot_labels`
    # has been treated for valid (non-ignore) pixels
    label, weight, valid_mask = _expand_onehot_labels(
        label, weight, pred.shape, ignore_index)
else:
    # should mask out the ignored elements
    valid_mask = ((label >= 0) & (label != ignore_index)).float()
    if weight is not None:
        weight = weight * valid_mask
    else:
        weight = valid_mask
# average loss over non-ignored and valid elements
if reduction == 'mean' and avg_factor is None and avg_non_ignore:
    avg_factor = valid_mask.sum().item()

loss = F.binary_cross_entropy_with_logits(
    pred, label.float(), pos_weight=class_weight, reduction='none')
# do the reduction for the weighted loss
loss = weight_reduce_loss(
    loss, weight, reduction=reduction, avg_factor=avg_factor)

return loss`

`def mask_cross_entropy(pred,
target,
label,
reduction='mean',
avg_factor=None,
class_weight=None,
ignore_index=None,
**kwargs):
"""Calculate the CrossEntropy loss for masks.

Args:
    pred (torch.Tensor): The prediction with shape (N, C), C is the number
        of classes.
    target (torch.Tensor): The learning label of the prediction.
    label (torch.Tensor): ``label`` indicates the class label of the mask'
        corresponding object. This will be used to select the mask in the
        of the class which the object belongs to when the mask prediction
        if not class-agnostic.
    reduction (str, optional): The method used to reduce the loss.
        Options are "none", "mean" and "sum".
    avg_factor (int, optional): Average factor that is used to average
        the loss. Defaults to None.
    class_weight (list[float], optional): The weight for each class.
    ignore_index (None): Placeholder, to be consistent with other loss.
        Default: None.

Returns:
    torch.Tensor: The calculated loss
"""
assert ignore_index is None, 'BCE loss does not support ignore_index'
# TODO: handle these two reserved arguments
assert reduction == 'mean' and avg_factor is None
num_rois = pred.size()[0]
inds = torch.arange(0, num_rois, dtype=torch.long, device=pred.device)
pred_slice = pred[inds, label].squeeze(1)
return F.binary_cross_entropy_with_logits(
    pred_slice, target, weight=class_weight, reduction='mean')[None]`