Training with threshold in PyTorch

learner47: <p>I have a neural network, which produces a single value when excited with input. I need to use this value returned by the network to threshold another array. The result of this threshold operation is used to compute a loss function (the value of threshold is not known before hand and needs to be arrived at by training).<br>
Following is an MWE</p>
<pre><code class="lang-auto">import torch

x = torch.randn(10, 1)  # Say this is the output of the network (10 is my batch size)
data_array = torch.randn(10, 2)  # This is the data I need to threshold
ground_truth = torch.randn(10, 2)  # This is the ground truth
mse_loss = torch.nn.MSELoss()  # Loss function

# Threshold
thresholded_vals = data_array * (data_array &gt;= x)  # Returns zero in all places where the value is less than the threshold, the value itself otherwise

# Compute loss and gradients
loss = mse_loss(thresholded_vals, ground_truth)
loss.backward()  # Throws error here
</code></pre>
<p>Since the operation of thresholding returns a tensor array that is devoid of any gradients the <code>backward()</code> operation throws error.</p>
<p>How does one train a network in such a case?</p>

suraj.pt: <p>Hi <a class="mention" href="/u/learner47">@learner47</a> I don’t think <code>x</code> can be differentiable here. Maybe try modifying the way you compute the threshold. Even then, your gradient is going to be 1 where <code>data_array&gt;=x</code> and 0 elsewhere. Maybe this discussion might help clarify things <a href="https://discuss.pytorch.org/t/how-to-make-the-parameter-of-torch-nn-threshold-learnable/4729/" class="inline-onebox">How to make the parameter of torch.nn.Threshold learnable?</a></p>
<p>Best!</p>

KFrank: <p>Hi Learner!</p>
<aside class="quote no-group quote-modified" data-username="learner47" data-post="1" data-topic="145962" data-full="true">
<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/l/5f9b8f/48.png" class="avatar"> learner47:</div>
<blockquote>
<pre><code class="lang-auto"># Threshold
thresholded_vals = data_array * (data_array &gt;= x)  # Returns zero in all places where the value is less than the threshold, the value itself otherwise
</code></pre>
<p>Since the operation of thresholding returns a tensor array that is devoid of any gradients the <code>backward()</code> operation throws error.</p>
</blockquote>
</aside>
<p>The thresholding operation is not (usefully) differentiable with respect<br>
to <code>x</code>.  To train <code>x</code> you should use a “soft,” differentiable thresholding<br>
operation.  You may use <code>sigmoid()</code> as a 'soft," differentiable step<br>
function.  Thus:</p>
<pre data-code-wrap="python"><code class="lang-python">thresholded_vals = data_array * torch.sigmoid (data_array - x)
</code></pre>
<p>You may introduce a parameter to sharpen or smooth such a “soft”<br>
step function:</p>
<pre data-code-wrap="python"><code class="lang-python">thresholded_vals = data_array * torch.sigmoid (alpha * (data_array - x))
</code></pre>
<p>As you increase <code>alpha</code> towards infinity, the thresholding sharpens into<br>
a hard step function.</p>
<p>Best.</p>
<p>K. Frank</p>