Retaining grad_fn for one-hot encoded tensors

stark: <p>Hi everyone,</p>
<p>I’m using <code>nn.functional.one_hot()</code> to convert predicted class labels <code>[N×1×H×W]</code> to one-hot encoded tensors <code>[N×C×H×W]</code>. The goal is to calculate loss by a custom loss function.</p>
<p>The problem is <code>nn.functional.one_hot()</code> sets <strong><code>grad_fn = None</code></strong>  so the calculated loss using such encoded predictions fails in backpropagation. How may I retain grad_fn for the one-hot encoded tensors?</p>
<p>Any thoughts are welcomed!</p>

KFrank: <p>Hi Stark!</p>
<aside class="quote no-group" data-username="stark" data-post="1" data-topic="123877" 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/stark/48/29013_2.png" class="avatar"> stark:</div>
<blockquote>
<p>How may I retain grad_fn for the one-hot encoded tensors?</p>
</blockquote>
</aside>
<p>You can’t, because it doesn’t make sense to.  <code>one_hot()</code> isn’t<br>
usefully differentiable, so a loss function that uses* it won’t be either.</p>
<p><code>one_hot()</code> takes a <code>torch.int64</code> argument and returns a <code>torch.int64</code><br>
result.  Pytorch doesn’t even permit such integer-valued tensors to<br>
have <code>requires_grad = True</code> (because it doesn’t make sense).</p>
<p>*)  To be more precise, a loss function could depend on the result<br>
of <code>one_hot()</code> and also on the results of some differentiable tensor<br>
operations.  Such a loss function would be usefully differentiable<br>
in that you could backpropagate through the differentiable tensor<br>
operations.  But you wouldn’t be able to backpropagate through the<br>
<code>one-hot()</code> part of the computation (nor anything upstream of it).</p>
<p>Best.</p>
<p>K. Frank</p>

stark: <p>Thanks for your reply.</p>
<p>What’re the options if one wants to use Dice or IOU as a loss function? Using class indices doesn’t make sense when an ordinal relationship doesn’t exist between classes. The only reasonable solution is to use one-hot encoding, but it’s useless at the cost of losing gradients. Any other ideas?</p>

KFrank: <p>Hi Stark!</p>
<aside class="quote no-group" data-username="stark" data-post="3" data-topic="123877" 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/stark/48/29013_2.png" class="avatar"> stark:</div>
<blockquote>
<p>What’re the options if one wants to use Dice or IOU as a loss function?</p>
</blockquote>
</aside>
<p>Let me imagine that your use case is something like multi-class<br>
semantic segmentation.  Your <code>target</code> is an “image” of class labels,<br>
that is this pixel is “class-3,” that pixel is “class-1,” and so on.  The<br>
prediction from your network is a set of <code>nClass</code> raw-score logits<br>
for each pixel, with shape <code>[nBatch, nClass, height, width]</code>.</p>
<p><code>Softmax</code> (perhaps better named <code>SoftOneHotEncodedArgMax</code>)<br>
will differentiably push your predictions towards <code>0.0</code> and <code>1.0</code>.</p>
<p>So use <code>soft_labels = torch.softmax (alpha * pred, dim = 1)</code></p>
<p><code>alpha</code> can be used to tune how “hard” your soft one-hot labels<br>
become – as you increase <code>alpha</code> you will push the labels closer<br>
to being exactly zero or one.  But don’t make <code>alpha</code> too large – if<br>
the labels saturate and become exactly zero and one you will lose<br>
useful differentiability, and as they get close to saturating, your<br>
differentiability will become less and less useful.</p>
<p>Probably the most straightforward way to calculate your “soft”<br>
intersection-over-union will be to one-hot your integer-class-label<br>
<code>target</code> so that it also has shape <code>[nBatch, nClass, height, width]</code>.<br>
Letting <code>pred</code> be your soft-label predictions and <code>targ</code> be your<br>
one-hotted class labels, then:</p>
<pre data-code-wrap="python"><code class="lang-python">iou = (pred * targ).sum (dim = (2, 3)) / (pred + targ).sum (dim = (2, 3))
</code></pre>
<p>(If <code>pred</code> and <code>targ</code> were boolean – that is just 0 or 1 – <code>pred * targ</code><br>
would be the boolean intersection and <code>pred + targ</code> the union.)</p>
<p><code>iou</code> will now have shape <code>[nBatch, nClass]</code> and each element of it<br>
will be the soft IOU for a specific class and a specific sample in the<br>
batch.</p>
<p>Best.</p>
<p>K. Frank</p>

stark: <p>Your explanation makes sense for <em>multi-class semantic segmentation</em>.</p>
<p>What if I want predicted class labels <code>[N×1×H×W]</code> from <code>torch.sigmoid(logits)</code> either as <code>0</code> or <code>1</code> and shape as <code>[N×C×H×W]</code> for a <strong>hard loss function</strong>? All such attempt <em>using thresholding, torch.one_hot(),…</em> fails to retain the gradient fucntion.</p>

KFrank: <p>Hi Stark!</p>
<aside class="quote no-group" data-username="stark" data-post="5" data-topic="123877" 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/stark/48/29013_2.png" class="avatar"> stark:</div>
<blockquote>
<p>Your explanation makes sense for <em>multi-class semantic segmentation</em>.</p>
<p>What if I want predicted class labels <code>[N×1×H×W]</code> from <code>torch.sigmoid(logits)</code> either as <code>0</code> or <code>1</code></p>
</blockquote>
</aside>
<p>If by this you mean that you wish to perform <em>binary</em> semantic<br>
segmentation, then you would do essentially as I described<br>
above, but use <code>torch.sigmoid (alpha * logits)</code> (rather than<br>
<code>torch.softmax()</code>).  And, analogously, the larger you make your<br>
choice of <code>alpha</code>, the “harder” your binary predictions will become.</p>
<aside class="quote no-group">
<blockquote>
<p>for a <strong>hard loss function</strong>? All such attempt <em>using thresholding, torch.one_hot(),…</em> fails to retain the gradient fucntion.</p>
</blockquote>
</aside>
<p>Yes, if you insist on a <strong>hard</strong> loss function, that is, something like<br>
<code>integer_valued_intersection / integer_valued_union</code>,<br>
you will, as you have seen, lose the gradients.</p>
<p>When you turn something that feeds into your loss function into an<br>
integer, you’ve applied a discontinuous function that isn’t usefully<br>
differentiable.  So you need to design a “soft” version of your desired<br>
loss function through which you can backpropagate.</p>
<p>As an aside, why not use <code>BCEWithLogitsLoss</code> (on a per-pixel basis)<br>
for your binary semantic segmentation?  People often use this loss<br>
function to good effect.</p>
<p>Best.</p>
<p>K. Frank</p>

stark: <p>Hi Frank,</p>
<aside class="quote no-group" data-username="KFrank" data-post="6" data-topic="123877">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>So you need to design a “soft” version of your desired<br>
loss function through which you can backpropagate.</p>
</blockquote>
</aside>
<p>That’s right. The use of <em>soft</em> predictions with <em>hard</em> targets in binary/semantic segmentation problems is commonly practiced.</p>
<p>Please share your thoughts on having a <strong><code>soft target</code></strong> in a soft loss function. Could it be a better option? I will be more than happy if you suggest some relevant work on this.</p>

KFrank: <p>Hi Stark!</p>
<aside class="quote no-group" data-username="stark" data-post="7" data-topic="123877" 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/stark/48/29013_2.png" class="avatar"> stark:</div>
<blockquote>
<p>Please share your thoughts on having a <strong><code>soft target</code></strong> in a soft loss function.</p>
</blockquote>
</aside>
<p>This is a perfectly reasonable thing to do. <code>BCEWithLogitsLoss</code> (as<br>
well as <code>BCELoss</code>) supports soft targets – you just pass them in as<br>
the <code>target</code>.  You can certainly contemplate using soft targets with<br>
other “soft” loss functions, with the details depending on the specific<br>
loss function.</p>
<p>(As an aside, it also makes sense to use soft labels with cross-entropy.<br>
Pytorch’s <code>CrossEntropyLoss</code> doesn’t support soft labels, but it is<br>
straightforward to write your own soft-label version of cross-entropy.)</p>
<aside class="quote no-group">
<blockquote>
<p>Could it be a better option?</p>
</blockquote>
</aside>
<p>Probably not (but it depends).</p>
<p>Your targets are what you <em>know</em> about your training data.  If you<br>
have (correct) hard class-“0” / class-“1” labels there is no benefit<br>
in artificially fuzzing them into soft labels – you would just be throwing<br>
away information.  On the other hand, if your labels are naturally<br>
soft, you probably make things worse by making them artificially<br>
hard.</p>
<p>To illustrate what I mean, consider training a binary cat-dog classifier.<br>
If all of your training images are clear cut – “Yup, that’s a cat; yup<br>
that’s a dog.” – then you should use hard labels because doing so<br>
uses all of the information you have.</p>
<p>But if some of your images are truly uncertain – maybe they are blurry<br>
or partially obscured – you should use soft labels, because that’s what<br>
you actually know.</p>
<p>In practice, why might you have soft labels and how would you get<br>
them?  You could, for example, have 100 people label each image.<br>
If 93 people label a specific image as “dog” (and 3 as “cat”), you<br>
would give that image of soft label of 0.93.  If 24 (of the 100) people<br>
label an image as “dog,” you would give it a soft label of  0.24.</p>
<p>The fact that people don’t always agree about the images indicates<br>
that the image labels are uncertain, that you might want to use soft<br>
labels, and gives you a quantitative way to assign such soft labels.</p>
<p>Best.</p>
<p>K. Frank</p>