Defining custom activation function

silvester: <p>I created a custom activation function MyReLU<br>
howevre when I use it in the two layer models I get the error</p>
<p><code> MyReLU.apply is not a Module subclass</code></p>
<p>MyReLU is a subclass of torch.autograd.Function</p>

googlebot: <p>You don’t use Function in places where Module is used, i.e. in <code>__init__</code> of main module. You just invoke MyReLU.apply in forward(). If you want to use Function is containers like nn.Sequential, you must wrap it in a Module.</p>

silvester: <p>When I use MyReLU.apply in forward of the method in a module, it does not work either:</p>
<pre><code class="lang-auto">class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = grad_output.clone()
        return grad_input


class weightedReLU(nn.Module):
    def __init__(self, weights = 1):
        super().__init__()
        self.weights = weights* nn.Parameter(torch.ones(1))
        self.weights = Variable(self.weights.data, requires_grad=True)
        
    def forward(self, input):
        ex = self.weights.cuda()*MyReLU.apply(input)
        return ex
</code></pre>

ecdrid: <p>Maybe <a href="https://discuss.pytorch.org/t/relu-with-leaky-derivative/32818/5">this</a> can help!</p>

silvester: <p>As Alex mentioned they are not working in containers like nn.sequential.</p>

googlebot: <p>IDK, your module works for me. If you use JIT, use additional wrapper:</p>
<pre><code class="lang-auto">@jit.ignore
def MyReLU(x):
    return MyReLU.apply(x)
</code></pre>

silvester: <p>Does it also work when you wrap the model insider nn.sequential?</p>
<p>what’s the difference between defining a class as autograd function of as nn.Module?</p>

googlebot: <aside class="quote no-group" data-username="silvester" data-post="7" data-topic="79913">
<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/s/6bbea6/48.png" class="avatar"> silvester:</div>
<blockquote>
<p>Does it also work when you wrap the model insider nn.sequential?</p>
</blockquote>
</aside>
<p>Yea, I think your error comes from some other place…</p>
<pre><code class="lang-auto">class M(nn.Module):
	def __init__(self):
		super(M, self).__init__()
		self.s = nn.Sequential(nn.ELU(), weightedReLU())
		
	
	def forward(self,x):
		return self.s(x)
	
m = M()
y=m.forward(torch.ones(5).requires_grad_())
y.sum().backward()
</code></pre>
<p>(your code untouched, except I removed cuda())</p>
<aside class="quote no-group" data-username="silvester" data-post="7" data-topic="79913">
<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/s/6bbea6/48.png" class="avatar"> silvester:</div>
<blockquote>
<p>what’s the difference between defining a class as autograd function of as nn.Module?</p>
</blockquote>
</aside>
<p>Function is not related to nn.Module, you don’t even create instances of it explicitly, it is just a way to provide backward(). Module base class provides has all usual facilities for submodules, parameters, module tree enumeration etc.</p>

silvester: <p>Thanks for your answer.</p>
<p>Suppose I need to include a learnable parameter in function autograd,<br>
similar to the function below</p>
<pre><code class="lang-auto">class CustAct(torch.nn.Module):
    def __init__(self, alphas=1):
        super(CustAct, self).__init__()
        self.alphas = alphas*torch.nn.Parameter(torch.ones(1))

        self.alphas.requiresGrad = True  

     def forward(self, x):
            val = self.alphas.cuda() * x

        return val
</code></pre>
<p>How I can include this in  autograd function?</p>
<p>Is there a way that I can remove cuda() so it can automatically differentiate between cpu and cuda? Even when I use CustAct.cuda(), still I need to explicitly map self.alphas to GPU.</p>

googlebot: <pre><code class="lang-auto">class CustAct(torch.nn.Module):
    def __init__(self, alphas=1):
        super(CustAct, self).__init__()
        self.alphas = torch.nn.Parameter(torch.full((1,), alphas))
        self.register_buffer("scale", torch.tensor([alphas]) #if you want to separate non-trainable constant, multiply in forward()
        self.scale = alphas #works too

     def forward(self, x):
        val = customAutogradFunction(x, self.alphas) #if you must provide custom gradients
        val = x * self.alphas #no need for autograd.Function
        val = x.clamp(min=0) #your MyReLU as is also needs no autograd.Function
        return val
</code></pre>
<aside class="quote no-group" data-username="silvester" data-post="9" data-topic="79913">
<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/s/6bbea6/48.png" class="avatar"> silvester:</div>
<blockquote>
<p>Is there a way that I can remove cuda() so it can automatically differentiate between cpu and cuda? Even when I use CustAct.cuda(), still I need to explicitly map self.alphas to GPU.</p>
</blockquote>
</aside>
<p>that’s because of the way you used multiplication, your self.alphas became Tensor not nn.Parameter and was not moved. Another idiom is tensorA2=tensorA.to(tensorB), this changes device and dtype, but it is not usually needed.</p>

silvester: <p>What is the effect of</p>
<aside class="quote no-group" data-username="googlebot" data-post="10" data-topic="79913">
<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/g/ba8739/48.png" class="avatar"> googlebot:</div>
<blockquote>
<p>self.register_buffer</p>
</blockquote>
</aside>
<p>in your code?<br>
I am loading some checkpoints and with this method, the loader asks for the values of alphas. It there any work around for this problem?<br>
How I could make a function in the way that it returns the gradients with respect to alphas, separately and not coupled with the other weights?</p>
<p>Is there a way that I can have separate alpha for each neuron? It seems that in this code, you used same self.alphas for all the elements of n dimensional array x.</p>

googlebot: <p>that code was purely illustrative, register_buffer is just a way to store non-trainable tensors in a module (sometimes it is more conventient even for scalars).</p>
<aside class="quote no-group" data-username="silvester" data-post="11" data-topic="79913">
<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/s/6bbea6/48.png" class="avatar"> silvester:</div>
<blockquote>
<p>How I could make a function in the way that it returns the gradients with respect to alphas, separately and not coupled with the other weights?</p>
</blockquote>
</aside>
<p>Why you want to manually return gradients in the first place? You usually only do this with external (thus non-differentiable) calculations, e.g. c++ extensions, or if you have simplified/more efficient gradient formulas. In other cases you only write forward() for nn.Module and most operations on tensors know how to calculate their gradients.</p>

silvester: <p>For this code, does it mean that weights are non-trainable any longer, as they changed to tensor?</p>
<pre><code class="lang-auto">class weightedReLU(nn.Module):
   def __init__(self, weights = 1): 
        super().__init__() 
        self.weights = weights* nn.Parameter(torch.ones(1)) 
        self.weights = Variable(self.weights.data, requires_grad=True) 
  def forward(self, input): 
       ex = self.weights.cuda()*MyReLU.apply(input) 
       return ex
</code></pre>

googlebot: <p>I don’t know what effects such a reassignment does, because that’s not an idiomatic pytorch code, to put it mildly. “Variable” is obsolete, and that line shouldn’t be there at all. For first assignment - yes, it may train, but your parameter won’t be registered.</p>

silvester: <p>I add the line for Variable, since otherwise it asked for<br>
<code>retain_graph=True</code><br>
for backward.</p>
<p>It seems Variable will register the gradient of weights in the backward method.</p>
<p>This code and yours produce completely different results.</p>

googlebot: <p>No, that’s probably error that you get if you do backward() twice.</p>
<aside class="quote no-group" data-username="silvester" data-post="15" data-topic="79913">
<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/s/6bbea6/48.png" class="avatar"> silvester:</div>
<blockquote>
<p>This code and yours produce completely different results.</p>
</blockquote>
</aside>
<p>Actually, weights* nn.Parameter(torch.ones(1)) won’t train as optimizer won’t find this parameter, so there is that. Oh, and then you may get that error. What a mess <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

silvester: <aside class="quote no-group" data-username="googlebot" data-post="16" data-topic="79913">
<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/g/ba8739/48.png" class="avatar"> googlebot:</div>
<blockquote>
<p>No, that’s probably error that you get if you do backward() twice.</p>
</blockquote>
</aside>
<p>My code with nn.ReLU works fine and with this custom activation after introducing Variable works fine as well (if I backward twice, it should be problematic with nn.ReLU too).</p>
<p>if you look at <a href="https://stackoverflow.com/questions/55765234/pytorch-custom-activation-functions" rel="noopener nofollow ugc">this post</a>,  it is suggested to do multiplication for the trainable variable.</p>

googlebot: <aside class="quote no-group" data-username="silvester" data-post="13" data-topic="79913">
<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/s/6bbea6/48.png" class="avatar"> silvester:</div>
<blockquote>
<p>self.weights = weights* nn.Parameter(torch.ones(1))</p>
</blockquote>
</aside>
<p>This creates untracked parameter - it is not in module, but in backward graph. Optimizer doesn’t process it, so you get error on second iteration.</p>
<pre><code class="lang-auto">self.weights = Variable(self.weights.data, requires_grad=True) 
</code></pre>
<p>this probably just makes first assignment have no effect, and is equivalent to</p>
<pre><code class="lang-auto">self.weights = nn.Parameter(torch.ones(1) * weights)
</code></pre>
<p>or torch.full((1,),weights)</p>
<aside class="quote no-group" data-username="silvester" data-post="17" data-topic="79913">
<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/s/6bbea6/48.png" class="avatar"> silvester:</div>
<blockquote>
<p>if you look at <a href="https://stackoverflow.com/questions/55765234/pytorch-custom-activation-functions" rel="noopener nofollow ugc">this post</a>, it is suggested to do multiplication for the trainable variable.</p>
</blockquote>
</aside>
<pre><code class="lang-auto">list(LearnedSwish().parameters())
[]
</code></pre>
<p>good luck optimizing that <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

silvester: <aside class="quote no-group" data-username="googlebot" data-post="18" data-topic="79913">
<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/g/ba8739/48.png" class="avatar"> googlebot:</div>
<blockquote>
<p>self.weights = Variable(self.weights.data, requires_grad=True)</p>
</blockquote>
</aside>
<p>I tried this and print the weights. It seems that they are not training, and stayed constant. Even when I load the check points, it does not ask for the values of those parameters across the layers.</p>
<p>For this one</p>
<p><code>self.weights = nn.Parameter(torch.ones(1) * weights)</code></p>
<p>is there a way that I impose the loader to ignore asking for the values in the case of loading from a check point with conventional nn.ReLU activation?</p>

googlebot: <p>not sure, there is nn.Module.load_state_dict(torch.load(PATH), strict=False), if you’re using training loop frameworks, check their docs…</p>