Custom loss function IoU is not differentiable. Can you create a differentiable IoU loss function for ML?

passiveradio: <p>Thank you for reading my post.<br>
I’m a college student, and currently developing the peak detection algorithm using CNN to determine the ideal convolution kernel which is representable as the ideal mother wavelet function that will maximize the peak detection accuracy.</p>
<p>I’ve tried to create my own IoU loss function for the CNN training model, but I failed.<br>
My own loss function is described as below.</p>
<pre><code class="lang-auto">
'''
1D intersection over union loss function class
'''
class IoU(nn.Module):
  def __init__(self, thresh: float = 0.5):
    super().__init__()
    self.thresh = thresh

  def forward(self, inputs: torch.Tensor, targets:torch.Tensor, weights: Optional[torch.Tensor] = None, smooth: float = 0.0) -&gt; Tensor:

    inputs = torch.where(inputs &lt; self.thresh, 0, 1)
    batch_size = targets.shape[0]

    intersect = torch.logical_and(inputs, targets)
    intersect = intersect.view(batch_size, -1).sum(-1)

    union = torch.logical_and(inputs, targets)
    union = union.view(batch_size, -1).sum(-1)

    IoU = (intersect + smooth) / (union + smooth)
    IoU = IoU.mean()

    return IoU
</code></pre>
<p>and I tried to test whether this works or not by this simple model described below.</p>
<pre><code class="lang-auto">
x = torch.tensor(0, 1001, 256) # e.g. [0, 200, 30, 1000, ...]
true = torch.tensor(0, 2, 256) # e.g. [0, 1, 1, 0, 1, ...]

model = nn.Linear(256, 256)
criterion = IoU()

output = model(x)
loss = criterion(output, true)
loss.backward() # I'm stuck on here, cause my loss func IoU is not differentiable

print(f"loss ={loss}")
print(f"model weight: {model.weight.grad}")
print(f"model params: [x.grad for x in {model.parameters()]")
</code></pre>
<p>And the output on the terminal is <code>RuntimeError: element 0 of variables does not require grad and does not have a grad_fn</code></p>
<p>This project is the time ever for me to use PyTorch, so I didn’t know what it meant at the first glance, but after my quick research, I figured out why this loss function fails (I’m not sure this is correct though)</p>
<h2>
<a name="my-loss-function-iou-is-not-differentiable-1" class="anchor" href="#my-loss-function-iou-is-not-differentiable-1"></a>my loss function IoU is not differentiable.</h2>
<p>and</p>
<h3>
<a name="this-is-where-the-chain-rule-of-this-loss-function-break-2" class="anchor" href="#this-is-where-the-chain-rule-of-this-loss-function-break-2"></a>This is where the chain rule of this loss function break.</h3>
<pre><code class="lang-auto">IoU = torch.nan_to_num(IoU)
IoU = IoU.mean()
</code></pre>
<p>Soon after I noticed this, I took a deeper look at the GitHub or stack overflow to find any other differentiable IoU loss function, but I’m still not sure <b>how to create a differentiable IoU loss function</b> (especially for 1D data).</p>
<p>If you have any experiences or insights around what I’m stuck on right now, please give me any instructions. Any advice will be welcomed and I hope that I can build my own IoU function for CNN model at the end.</p>
<p>Thank you</p>

KFrank: <p>Hi Passive!</p>
<aside class="quote no-group quote-modified" data-username="passiveradio" data-post="1" data-topic="148899" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/passiveradio/48/48543_2.png" class="avatar"> passiveradio:</div>
<blockquote>
<pre><code class="lang-auto">  def forward(self, inputs: torch.Tensor, targets:torch.Tensor, weights: Optional[torch.Tensor] = None, smooth: float = 0.0) -&gt; Tensor:

    inputs = torch.where(inputs &lt; self.thresh, 0, 1)
</code></pre>
<h2>my loss function IoU is not differentiable.</h2>
</blockquote>
</aside>
<p>As you have deduced, your <code>IoU</code> loss criterion is not (usefully) differentiable.</p>
<p>First, to clear up some context:  Based on your thresholding value of<br>
<code>thresh = 0.5</code>, let me assume that <code>inputs</code> (the <code>output</code> of your <code>model</code>)<br>
are to be thought of as probabilities that range from zero to one (and<br>
represent the probability of the “pixel” in question being in the predicted<br>
“object” whose intersection-over-union you wish to calculate).  Note, this<br>
is not consistent with your <code>model</code> being just a <code>Linear</code> whose outputs<br>
range from <code>-inf</code> to <code>+inf</code> and would be thought of as the <em>logits</em> that<br>
correspond to the aforementioned probabilities.  If you have such logits<br>
in your real use case, you can convert them to probabilities by passing<br>
them through <code>sigmoid()</code>.</p>
<p>You threshold your probabilities to produce “hard,” yes-no predictions for<br>
whether a given pixel is predicted to be in the object, and then use these<br>
hard predictions in your <code>IoU</code> computation.  Although mathematically<br>
differentiable almost everywhere (everywhere except <code>x = thresh</code>, where<br>
the derivative is undefined), the derivative is always zero, so it gives your<br>
optimization no useful information.  (As an aside, as written, you cannot<br>
backpropagate through your thresholding.  But even if you could, you would<br>
just get zero.)</p>
<p>Calculate instead a “soft,” probabilistic <code>IoU</code>.  Replace your <code>logical_and()</code><br>
with multiplication:</p>
<pre data-code-wrap="python"><code class="lang-python">intersect = (inputs * targets).view (batch_size, -1).sum (-1)
</code></pre>
<p>(<code>logical_and()</code> is precisely the regular product when all the values are<br>
either zero or one.)</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">    union = torch.logical_and(inputs, targets)
</code></pre>
</blockquote>
</aside>
<p>Note, you most likely have a typo in your expression for <code>union</code> and<br>
presumably meant <code>logical_or()</code>.</p>
<p>Off the top of my head, I would replace the <code>logical_or()</code> that you (should)<br>
have in the expression for <code>union</code> with <code>max</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">intersect = torch.max (inputs, targets).view (batch_size, -1).sum (-1)
</code></pre>
<p>(Again, <code>max()</code> is the same as <code>logical_or()</code> when all of the values are<br>
zero or one.)</p>
<p>If you want to threshold <code>inputs</code>, rather than just use the “raw” probabilities,<br>
you can use a “soft,” differentiable threshold:</p>
<pre><code class="lang-auto">thresholded_inputs = inputs**alpha / (inputs**alpha + (1 - inputs)**alpha)
</code></pre>
<p>(with <code>alpha &gt;= 0.0</code>).</p>
<p><code>alpha</code> is a parameter that “sharpens” the thresholding:  As <code>alpha</code> ranges<br>
from <code>1.0</code> to <code>+inf</code>, you interpolate between using the raw, unthresholded<br>
probabilities to fully thresholding them to be exactly either zero or one.</p>
<p>(If you are predicting logits, rather than probabilities, just multiply the<br>
logits by <code>alpha</code> before passing them through <code>sigmoid()</code> to convert<br>
them to probabilities.)</p>
<p>(As an aside, using <code>IoU</code> suggests that you are trying to predict which<br>
“pixels” are in certain “objects.”  This is a binary segmentation problem<br>
problem, which is to say a per-pixel binary classification problem.  You<br>
might consider using <code>BCEWithLogitsLoss</code> as your loss function, with<br>
<code>pos_weight</code>, if necessary, to compensate for foreground-background<br>
class imbalance.  If you feel that using <code>IoU</code> would add value as the loss<br>
function, you might consider <em>augmenting</em> <code>BCEWithLogitsLoss</code> with <code>IoU</code>,<br>
or (a differentiable version of) its cousin, Dice loss.)</p>
<p>Good luck.</p>
<p>K. Frank</p>

passiveradio: <p>Hi! K. Frank<br>
Thank you for your decent advice, it helps a lot.<br>
I should’ve remembered the operation <b>matrix * matrix</b> doesn’t break the chain rule of differentiation <img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=12" title=":joy:" class="emoji" alt=":joy:" loading="lazy" width="20" height="20"><br>
I rewrote my model and loss function to make them differentiable and ready to run state according to your series of advice, and I succeeded to “execute” my training without an error!<br>
This is a huge advancement for me (plz allow me to say thank you again), though, I faced another problem (it might be a easy peasy to fix for knowledgeable person though) which is,</p>
<h2>
<a name="the-weights-of-the-model-doesnt-upgraded-through-its-learning-over-epochs-1" class="anchor" href="#the-weights-of-the-model-doesnt-upgraded-through-its-learning-over-epochs-1"></a>the weights of the model doesn’t upgraded through its learning over epochs.</h2>
<p>This must be another topic, so I posted a new topic <a href="https://discuss.pytorch.org/t/the-parameters-of-the-model-with-custom-loss-function-doesnt-upgraded-thorough-its-learning-over-epochs/149024">The parameters of the model with custom loss function doesn’t upgraded thorough its learning over epochs - discuss.pytorch.org</a></p>
<p>Thank you!<br>
passiveradio</p>