Can you apply transformations in Custom Loss functions

1

I have an idea for a Custom IoU Loss function. Which I implemented like shown below. I get

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

So my question is can I actually perform these transformation on my output and compute a loss value, or am I breaking fundamental chain rules. Explanation with reasoning or sources would be lovely. Thanks

class SobelIoULoss(nn.Module):
    def __init__(self, sobel_threshold):
        super(SobelIoULoss, self).__init__()
        self.sobel_threshold = sobel_threshold
        self.sobel_kernel = torch.tensor([[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]]],
                                          dtype=torch.float32).to('cuda')

    def forward(self, true, pred):

        # take the cube of inputs
        true = true**3
        pred = pred**3
        true.requires_grad_()
        pred.requires_grad_()

        true_sobel = F.conv2d(true, self.sobel_kernel, padding=1)
        pred_sobel = F.conv2d(pred, self.sobel_kernel, padding=1)

        # Create binary masks for true and pred
        true_mask = (true_sobel > self.sobel_threshold).float()
        pred_mask = (pred_sobel > self.sobel_threshold).float()

        # Compute intersection and union
        intersection = torch.sum(true_mask * pred_mask, dim=(1, 2, 3))
        union = torch.sum(true_mask + pred_mask - true_mask * pred_mask, dim=(1, 2, 3))

        # Compute IoU
        iou = intersection / union

        # Compute loss as 1 - IoU
        loss = 1 - iou

        # Return average loss of the batch
        return torch.mean(loss)
2

The comparison operation:

        # Create binary masks for true and pred
        true_mask = (true_sobel > self.sobel_threshold).float()
        pred_mask = (pred_sobel > self.sobel_threshold).float()

is detaching the tensor from the computation graph as it’s not usefully differentiable. The gradient would be zero almost everywhere and undefined or Inf at the rounding point.

3

Is there a workaround to achieve my idea? Basically I want to use IoU loss when not dealing with rectangles where you can estimate

x1, x2, y1, y2.

4

You could try to use a “soft” threshold as described here.