BCEWithLogitsLoss can't compare Segmentation Mask with predicted mask

Sowmen_Das: <p>I am doing a binary segmentation task. I’m using BCEWithLogitsLoss. I’m using Albumentations to augment and normalize images.</p>
<pre><code class="lang-auto">transforms_normalize = albumentations.Compose(
        [
            albumentations.Normalize(mean=normalize['mean'], std=normalize['std'], always_apply=True, p=1),
            albumentations.pytorch.transforms.ToTensorV2()
        ],
        additional_targets={'ela':'image'}
    )
</code></pre>
<p>This loads two images and a segmenattion mask. The images are returned as FloatTensor and mask is returned as ByteTensor.</p>
<p>When I try to calculate loss the following error is thrown</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/media/sowmitra/SSD Disk/image_manipulation/train_segment.py", line 559, in &lt;module&gt;
    resume=False
  File "/media/sowmitra/SSD Disk/image_manipulation/train_segment.py", line 243, in train
    train_metrics = train_epoch(model, train_loader, optimizer, criterion, epoch, SRM_FLAG)
  File "/media/sowmitra/SSD Disk/image_manipulation/train_segment.py", line 404, in train_epoch
    loss_segmentation = criterion(out_mask, gt)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/pytorch_toolbelt/losses/joint_loss.py", line 32, in forward
    return self.first(*input) + self.second(*input)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/pytorch_toolbelt/losses/joint_loss.py", line 18, in forward
    return self.loss(*input) * self.weight
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 632, in forward
    reduction=self.reduction)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/functional.py", line 2580, in binary_cross_entropy_with_logits
    raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
ValueError: Target size (torch.Size([8, 256, 256])) must be the same as input size (torch.Size([8, 1, 256, 256]))
</code></pre>
<p>To fix this I used <code>squeeze(1)</code> on the output tensor to remove the second dim. Should I squeeze the output or should I expand the ground-truth mask?</p>
<p>After doing that, this next error is thrown:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/media/sowmitra/SSD Disk/image_manipulation/train_segment.py", line 559, in &lt;module&gt;
    resume=False
  File "/media/sowmitra/SSD Disk/image_manipulation/train_segment.py", line 243, in train
    train_metrics = train_epoch(model, train_loader, optimizer, criterion, epoch, SRM_FLAG)
  File "/media/sowmitra/SSD Disk/image_manipulation/train_segment.py", line 404, in train_epoch
    loss_segmentation = criterion(out_mask, gt)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/pytorch_toolbelt/losses/joint_loss.py", line 32, in forward
    return self.first(*input) + self.second(*input)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/pytorch_toolbelt/losses/joint_loss.py", line 18, in forward
    return self.loss(*input) * self.weight
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 632, in forward
    reduction=self.reduction)
  File "/home/sowmitra/anaconda3/envs/dfdcpy37env/lib/python3.7/site-packages/torch/nn/functional.py", line 2582, in binary_cross_entropy_with_logits
    return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
RuntimeError: result type Float can't be cast to the desired output type Byte
</code></pre>
<p>I’m guessing this is because the model returns logits of size (nbatch,1,256,256) as FloatTensors since they are probabilities but the masks are (nbatch,256,256) ByteTensors. Should I just convert the masks to Float? Or should I threshold the outputs and convert it to Byte?</p>

ptrblck: <p>For a multi-class segmentation use case, you should use <code>nn.CrossEntropyLoss</code> where the model output should be a <code>FloatTensor</code> in the shape <code>[batch_size, nb_classes, height, width]</code> and the target a <code>LongTensor</code> in the shape <code>[batch_size, height, width]</code> containing the class indices in the range <code>[0, nb_classes-1]</code>.<br>
If you are dealing with a binary or multi-label segmentation, you could use <code>nn.BCEWithLogitsLoss</code>, which would expect a target as a <code>FloatTensor</code> in the same shape as the model output containing values in <code>[0, 1]</code>.</p>