This is an example of the output for a certain point:
image2083×1036 56 KB
On the left I have the ground truth for one of the channels, a dot formed by a Gaussian with one pixel of variance. On the right the outputted map on epoch 49. My dataset has around 20k examples. By epoch 49 it should have seen enough examples.
Here is my training step:
   def training_step(self, batch, batch_idx):

       x, mask = batch
       mask_pred = self(x)
       mask_pred = torch.sigmoid(mask_pred)
       loss = F.mse_loss(mask_pred, mask)

       self.log(
           "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True
       )

       if batch_idx % 200 == 0:
           self.log_masks(mask, mask_pred, batch_idx)

       return loss

The Gaussian is normalized around 0 - 1 so I apply Sigmoid to the output.