Learnable threshold parameter inside a function

Jg-lee: <p>Hi, I want to define custom layer with a learnable parameter <strong><code>T</code></strong>.<br>
The function behaves like below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/d/fd69e3a54da6fe9e02a398ae8513719d93ff36c8.png" data-download-href="https://discuss.pytorch.org/uploads/default/fd69e3a54da6fe9e02a398ae8513719d93ff36c8" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd69e3a54da6fe9e02a398ae8513719d93ff36c8_2_517x98.png" alt="image" data-base62-sha1="A9NFkHDmlhOM8ha6SEhApQ9xysM" width="517" height="98" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd69e3a54da6fe9e02a398ae8513719d93ff36c8_2_517x98.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd69e3a54da6fe9e02a398ae8513719d93ff36c8_2_775x147.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd69e3a54da6fe9e02a398ae8513719d93ff36c8_2_1034x196.png 2x" data-dominant-color="F4F4F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1108×211 10.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>here, x is the layer’s input, and T is a learnable parameter.<br>
I am thinking of defining my own autograd function something like below</p>
<pre data-code-wrap="python"><code class="lang-python">class _NewActivation(torch.autograd.Function):
   @staticmethod
   def forward(ctx, x, T):
      ctx.save_for_backward(x,T)
      # ret = f(x) for x&gt;=T, g(x) for x &lt;T
      # return ret
   
   @staticmethod
   def backward(ctx, grad_output):
      x, T = ctx.saved_tensors
      # compute_grad will compute grad_output * d_Activation/d_input
      return compute_grad(x, T, grad_output) 
</code></pre>
<p>And then add above function into my model by wrapping inside torch.nn.Module somthing like below.</p>
<pre data-code-wrap="python"><code class="lang-python">class NewActivation(torch.nn.Module):
   def __init__(self):
      super().__init__()
      self.T=Parameter(torch.Tensor([[0]]))
      self.activation = _NewActivation.apply
   def forward(self, x):
      return activation(x)
</code></pre>
<p>My question is, I don’t know how to make the <strong><code>T</code></strong> learnable. I think I shoud add another torch.autograd.Function to compute gradient of <code>d_Activation/d_T</code>, but I have no idea.</p>

Jg-lee: <p>Seems like we can match input of <strong>forward</strong> function with output gradient of <strong>backward</strong> function</p>
<pre><code class="lang-python">class _NewActivation(torch.autograd.Function):
   @staticmethod
   def forward(ctx, x, T):
      ctx.save_for_backward(x,T)
      return compute_ret(x,T)
   
   @staticmethod
   def backward(ctx, grad_output):
      x, T = ctx.saved_tensors
      x_grad = compute_x_grad(x, T, grad_output) # d_Loss/dx
      T_grad = compute_T_grad(x, T, grad_output) # d_Loss/dT
      return x_grad, T_grad # Position should match with forward's positional argument 
</code></pre>
<p>I tried simple test cases and think this is valid to set gradient of x, T (input and parameters of the layer). However, I’m not sure this is widely used way for implementing custom autograd function with learnable parameter…</p>

ptrblck: <p>Your custom <code>autograd.Function</code> pseudo-code looks alright, but I wouldn’t know how you want to calculate the gradient of <code>T</code> as it should be zero almost everywhere and undefined/Inf at <code>T==x</code>, shouldn’t it?</p>

Jg-lee: <p>Oh, actually I have some T terms in the f, g like:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/f/0f75ee7ed3654e619464f7123d675742a2ee54b1.png" data-download-href="https://discuss.pytorch.org/uploads/default/0f75ee7ed3654e619464f7123d675742a2ee54b1" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f75ee7ed3654e619464f7123d675742a2ee54b1_2_690x74.png" alt="image" data-base62-sha1="2cLPqbB5iKjGxo7fBXADaV2KmMF" width="690" height="74" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f75ee7ed3654e619464f7123d675742a2ee54b1_2_690x74.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f75ee7ed3654e619464f7123d675742a2ee54b1_2_1035x111.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f75ee7ed3654e619464f7123d675742a2ee54b1_2_1380x148.png 2x" data-dominant-color="F1F1F1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2195×238 21.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
I can calculate my gradient of T by fixing x to input value.</p>