Problems in merging customized neuron into standard ones

username: bestcql
cooked: <p>I created a customized autograd function and wrap it inside a module. But I just keep getting errors when I try to merge my module with other standard modules using nn.Sequential(). Can anyone help? Thanks.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

# customized autograd function
class BinaryTanh_func(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        tanh_func = nn.Tanh()
        preactivated_threshold = tanh_func(input)
        final_output = torch.ones(preactivated_threshold.size())
        final_output[preactivated_threshold &lt; 0] = 0
        return final_output

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        tanh_func = nn.Tanh()
        preactivated = tanh_func(input)
        grad_input = (1-preactivated**2) * grad_input
        return grad_input

# customized module
class BinaryTanh(nn.Module):
    @staticmethod
    def forward(self, input):
        return BinaryTanh_func()(input)


# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

torch.manual_seed(100)
# Create random Tensors to hold inputs and outputs
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    BinaryTanh(),
    torch.nn.Linear(H, D_out),
)

# The nn package also contains definitions of popular loss functions; in this
# case we will use Mean Squared Error (MSE) as our loss function.
loss_fn = torch.nn.MSELoss(reduction='sum')

learning_rate = 1e-4
for t in range(10):
    y_pred = model(x)
    loss = loss_fn(y_pred, y)
    print(t, loss.item())
    model.zero_grad()
    loss.backward()

    with torch.no_grad():
        for param in model.parameters():
            param -= learning_rate * param.grad
</code></pre>
<h2>Here’s the error message I got:</h2>
<p>TypeError                                 Traceback (most recent call last)</p>
<pre><code class="lang-auto">&lt;ipython-input-29-becf74c80a68&gt; in &lt;module&gt;()
     73 for t in range(10):
     74 
---&gt; 75     y_pred = model(x)
     76 
     77     loss = loss_fn(y_pred, y)

2 frames

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    491             result = self._slow_forward(*input, **kwargs)
    492         else:
--&gt; 493             result = self.forward(*input, **kwargs)
    494         for hook in self._forward_hooks.values():
    495             hook_result = hook(self, input, result)

TypeError: forward() missing 1 required positional argument: 'input'
</code></pre>
<hr>
<p>Any help will be appreciated. Thanks.</p>

username: ptrblck
cooked: <p>Try to change your <code>BinaryTanh</code> definition to:</p>
<pre><code class="lang-python"># customized module
class BinaryTanh(nn.Module):
    def __init__(self):
        super().__init__()
        self.fn = BinaryTanh_func.apply
        
    def forward(self, input):
        return self.fn(input)
</code></pre>

username: bestcql
cooked: <p>Thanks a lot!! It works.<br>
And I also found that if I comment the ‘’<span class="mention">@staticmethod</span>" line in my original code, it will work too. Can you briefly explain the reason? I think I will dig deeper into this and try to google more about the “staticmethod” thing. Thanks again.</p>

username: ptrblck
cooked: <p><code>staticmethods</code> are bound to the class and do not require a class instance. They are therefore independent of the “state” of a class instance.<br>
So instead of creating an instance (attributes) and calling it, you just might call <code>MyFunction.apply</code> directly.</p>

username: bestcql
cooked: <p>Got it. Thank you very much!</p>