Quantized hard sigmoid

username: pshashk
cooked: <p>I’ve tried to implement hard sigmoid activation in a way suitable for quantization aware training:</p>
<pre><code class="lang-auto">from torch import nn

class HardSigmoid(nn.Module):
    def __init__(self):
        super().__init__()
        self.act = nn.ReLU6()
        self.add = nn.quantized.FloatFunctional()
        self.mul = nn.quantized.FloatFunctional()
    
    def forward(self, input):
        # relu6(input + 3) / 6
        output = self.add.add_scalar(input, 3)
        output = self.act(output)
        output = self.mul.mul_scalar(output, 1/6)
        return output
</code></pre>
<p>The backward pass and conversion works fine:</p>
<pre><code class="lang-auto">import torch.quantization as tq
from torch.nn.intrinsic import ConvBn2d
import torch

model = nn.Sequential(
    tq.QuantStub(), 
    ConvBn2d(
        nn.Conv2d(3, 16, kernel_size=3, padding = 1, bias = False),
        nn.BatchNorm2d(16)
    ),
    HardSigmoid(),
    tq.DeQuantStub()
)

optimizer = torch.optim.SGD(model.parameters(), lr = 0.0001)
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace = True)

x = torch.rand(16, 3, 16, 16)
y = model(x)

y.sum().backward()
optimizer.step()

y = model(x)

model.apply(torch.quantization.disable_observer)
model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)

torch.quantization.convert(model.eval(), inplace = True)
</code></pre>
<p>But the forward pass (<code>model(x)</code>) fails with the following error:</p>
<pre><code class="lang-auto">  File "/usr/local/lib64/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib64/python3.7/site-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/usr/local/lib64/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "&lt;stdin&gt;", line 10, in forward
  File "/usr/local/lib64/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib64/python3.7/site-packages/torch/nn/modules/activation.py", line 209, in forward
    return F.hardtanh(input, self.min_val, self.max_val, self.inplace)
  File "/usr/local/lib64/python3.7/site-packages/torch/nn/functional.py", line 960, in hardtanh
    result = torch._C._nn.hardtanh(input, min_val, max_val)
RuntimeError: Didn't find kernel to dispatch to for operator 'aten::hardtanh'. Tried to look up kernel for dispatch key 'QuantizedCPUTensorId'. Registered dispatch keys are: [CUDATensorId, CPUTensorId, VariableTensorId]
</code></pre>
<p>What is the correct way to implement hard sigmoid activation?</p>

username: jerryzh168
cooked: <aside class="quote no-group" data-username="pshashk" data-post="1" data-topic="59013">
<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/p/9d8465/48.png" class="avatar"> pshashk:</div>
<blockquote>
<p>tq.DeQuantStub()</p>
</blockquote>
</aside>
<p>We actually don’t have relu6 in the mapping right now: <a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/default_mappings.py#L14" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/torch/quantization/default_mappings.py#L14</a>. adding an entry here should fix the problem. Thanks for reporting.</p>