Gradient flow through torch.count_nonzero

cbd: <p>As per below code the, the backward function is called after run. But once i uncomment the line “output=torch.count_nonzero(output)” it shows the error “RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn”. It is necessary to use “torch.count_nonzero” function and wants to pass the “grad_input” 1 as it is non differentiable.</p>
<pre><code class="lang-python">from torch.autograd import Function
import torch
import torch.nn as nn

class custom_fun(Function):

    @staticmethod
    def forward(cxt, input):
        cxt.save_for_backward(input)
        output=torch.add(input,2)
        print("calling forward",output)
        # output=torch.count_nonzero(output)
        return output

    @staticmethod
    def backward(cxt, grad_output):
        input, = cxt.saved_tensors
        print("Calling backward function ")
        grad_input = grad_output.clone()
        grad_input = torch.ones(1)
        return grad_input
class my_module(nn.Module):

 def __init__(self):
      super(my_module,self).__init__()
      self.fc1 = nn.Linear(1,1)
      self.my_act_fn=custom_fun()
  def forward(self, x):
      out1=self.fc1(x)
      out=self.my_act_fn.apply(out1)
      return out

mod = my_module()
x=torch.tensor([1.1])
target=torch.tensor([2.5])
out_m=mod(x)
criterion=torch.nn.L1Loss()
loss=criterion(out_m,target)
loss.backward()
</code></pre>

albanD: <p>Hi,</p>
<p>The problem is that the output of <code>count_nonzero</code> is a Tensor of type int64. So it cannot require gradients (only continuous types can).<br>
So you want to add a <code>output = output.float()</code> after the count_nonzero to make sure the output of your custom Function can require gradients.</p>

cbd: <p>Thanks. Now gradient flows.</p>

sgbaird: <p>Is <a class="mention" href="/u/cbd">@cbd</a>’s custom function a sound approach? As seems to be mentioned by OP, <code>torch.count_nonzero</code> is non-differentiable. If I understand correctly, is this essentially ignoring the gradient of <code>torch.count_nonzero</code>?</p>

sgbaird: <p>Rather than use <code>torch.count_nonzero</code>, which is non-differentiable, an alternative is to approximate this via a “differentiable relaxation using a sum of narrow Gaussian basis functions”. i.e.</p>
<pre><code class="lang-python">def narrow_gaussian(x, ell):
    return torch.exp(-0.5 * (x / ell) ** 2)

def approx_count_nonzero(x, ell=1e-3):
    # Approximation of || x ||_0
    return len(x) - narrow_gaussian(x, ell).sum(dim=-1)
</code></pre>
<p>As the value for the i-th component (<code>x[i]</code> or x<sub>i</sub>) approaches zero, the Gaussian basis function evaluated at <code>x</code> tends towards <code>1</code>:</p>
<p>lim<sub>x→0</sub> f(x<sub>i</sub>) = 1</p>
<p>Likewise, as the value of the i-th component deviates from zero, the Gaussian basis function evaluated at <code>x</code> tends towards <code>0</code>:</p>
<p>lim<sub>x→∞</sub> f(x<sub>i</sub>) = 0</p>
<p>So if you have three components with values close to zero and the rest not close to zero, then the sum of all the narrow Gaussians is approximately <code>3</code>, and the number of non-zero components is <code>len(x)</code> minus this sum.</p>
<p>modified from <a href="https://github.com/facebook/Ax/files/7936069/Scipy.non-linear.inequality.constraints.txt" rel="noopener nofollow ugc">Scipy non-linear inequality constraints.ipynb</a> as described in <a href="https://github.com/facebook/Ax/issues/769#issuecomment-1021471675" class="inline-onebox" rel="noopener nofollow ugc">[FEATURE REQUEST] modify Ax API to allow for callable that evaluates a constraint and is passed to the optimizer · Issue #769 · facebook/Ax · GitHub</a>.</p>
<p>Credit: David Eriksson, Senior Research Scientist at Meta</p>