Binary Activation Function with Pytorch

Arhazf: <p>I have a 2-layers fully connected network. I would like to convert the output of the first layer to binary. This means that I would like to have a binary-step activation function in the forward paths and Relu activation function in the backward pass. How can I implement this? Any idea would be appreciated.</p>
<pre><code>class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 64)
        self.fc2 = nn.Linear(64, 10)
        nn.init.xavier_normal_(self.fc1.weight)
        nn.init.xavier_normal_(self.fc2.weight)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        x = F.log_softmax(x, dim=1)
        return x

net = Net()</code></pre>

ptrblck: <p><a href="https://discuss.pytorch.org/t/relu-with-leaky-derivative/32818/2">This example</a> could probably do, what you need.</p>

Arhazf: <p>Thank you for your point. I rewrote my code according to the example:</p>
<pre><code>class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 64)
        self.fc2 = nn.Linear(64, 10)
        nn.init.xavier_normal_(self.fc1.weight)
        nn.init.xavier_normal_(self.fc2.weight)

    def forward(self, x):
        x1 = F.relu(self.fc1(x))
        x_backward = x1
        x1[x1 &lt;= 0] = 0
        x1[x1&gt; 0] = 1
        x_forward = x1
        y1 = x_backward + (x_forward - x_backward).detach()
        y2 = self.fc2(y1)
        y3 = F.log_softmax(y2, dim=1)
        return y3  
</code></pre>
<p>But it does not work and it returns binary values during both forward and backward phases.</p>
<p>I also tried to define a new class and define the forward and backward paths separately…</p>
<p>I tried to do it like this but it gives me an error (AttributeError: ‘Binary_AF’ object has no attribute ‘dim’ ).</p>
<p>Do you think it makes sense to write the code like this? if yes, how to take care of the error and if no, I would appreciate if you have any suggestion.</p>
<pre><code>class Binary_AF:
    def __init__(self, x):
        self.x = x

    def forward(self):
        self.x[self.x &lt;= 0] = 0
        self.x[self.x &gt; 0] = 1
        return self.x

    def backward(self):
        return self.x


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 64)
        self.fc2 = nn.Linear(64, 10)
        nn.init.xavier_normal_(self.fc1.weight)
        nn.init.xavier_normal_(self.fc2.weight)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        y = Binary_AF(x)
        y = self.fc2(y)
        y = F.log_softmax(y, dim=1)
        return y</code></pre>

wayne391: <p>Hi, I think there is a simper way by combining ‘sign’ and ‘relu’:</p>
<pre data-code-wrap="python"><code class="lang-python"># test input
input = torch.tensor([-20, 11], requires_grad=True).float()

# binary
sign = torch.sign(input)
binary_out = torch.relu(sign)
print(binary_out) # 0., 1.

</code></pre>
<p>Theoretically, sign (or unit-step) is not differentiable, but it can be backpropagated in PyTorch.<br>
I have no idea how they design it…</p>
<p>Here is a question about this:</p><aside class="quote quote-modified" data-post="1" data-topic="41718">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/soshishimada/48/2278_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/circumventing-an-indifferentiablility-problem/41718">Circumventing an indifferentiablility problem</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Sorry if my question isn’t appropriate to ask here! It’s a bit theory related question. But I want people to share knowledge to understand what’s happening in neural networks. 
As far as I understand, Pytorch use chain rule to compute gradients of loss w.r.t. network parameters. 
Therefore, when we use an indifferentiable function such as step function (torch.sign() ) in the neural network, the gradient won’t be propagated hence loss won’t decrease. 
In the code below, I implemented a very simpl…
  </blockquote>
</aside>


Filip_Cornell: <p>For this, for future visitors, I can recommend <a href="https://www.hassanaskary.com/python/pytorch/deep%20learning/2020/09/19/intuitive-explanation-of-straight-through-estimators.html#:~:text=A%20straight%2Dthrough%20estimator%20is,function%20was%20an%20identity%20function." rel="noopener nofollow ugc">this blogpost</a> explaining the <a href="https://arxiv.org/pdf/1308.3432.pdf" rel="noopener nofollow ugc">Straight-Through Estimator</a> (Bengio et al.).</p>
<p>For a quick fix, this module (directly copied from the blog post linked above) should be able to handle a binary activation function.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class STEFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        return (input &gt; 0).float()

    @staticmethod
    def backward(ctx, grad_output):
        return F.hardtanh(grad_output)

class StraightThroughEstimator(nn.Module):
    def __init__(self):
        super(StraightThroughEstimator, self).__init__()

    def forward(self, x):
        x = STEFunction.apply(x)
        return x

x = torch.randn(30,30)
print(x)
estimator = StraightThroughEstimator()
estimator(x)
# tensor([[-1.6039e+00, -3.1636e-01, -1.4391e+00,  8.4508e-02, -2.6198e+00, ...
# tensor([[0., 0., 0., 1., 0., 1., 0., 1., 0., 1., 1., 0., 1., 0., 1., 0., 0., 0., ...
</code></pre>