Creating a custom layer and using torch.qat for it

username: thyeros
cooked: <p>Hi, I have a specific case and need some help/pointer.</p>
<p>I have designed a specialized normalization layer (with learnable parms) from nn.Module and like to apply QAT to this one. But, I couldn’t find a doc on how to make a corresponding module for QAT: such as attaching processes weight_fake_quant, activation_post_process as well.</p>
<p>Any starter on this?</p>

username: jerryzh168
cooked: <p>we have examples for qat.Conv2d here: <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/qat/modules/conv.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch/conv.py at master · pytorch/pytorch · GitHub</a></p>
<p>then you can add the mapping in:</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantization_mappings.py#L68" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantization_mappings.py#L68" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/master/torch/quantization/quantization_mappings.py#L68</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="58" style="counter-reset: li-counter 57 ;">
<li>    nniqat.ConvBnReLU1d: nniq.ConvReLU1d,</li><li>    nniqat.ConvBnReLU2d: nniq.ConvReLU2d,</li><li>    nniqat.ConvReLU2d: nniq.ConvReLU2d,</li><li>    nniqat.LinearReLU: nniq.LinearReLU,</li><li>    # QAT modules:</li><li>    nnqat.Linear: nnq.Linear,</li><li>    nnqat.Conv2d: nnq.Conv2d,</li><li>}</li><li></li><li># Default map for swapping float module to qat modules</li><li class="selected">DEFAULT_QAT_MODULE_MAPPINGS : Dict[Callable, Any] = {</li><li>    nn.Conv2d: nnqat.Conv2d,</li><li>    nn.Linear: nnqat.Linear,</li><li>    nn.modules.linear._LinearWithBias: nnqat.Linear,</li><li>    # Intrinsic modules:</li><li>    nni.ConvBn1d: nniqat.ConvBn1d,</li><li>    nni.ConvBn2d: nniqat.ConvBn2d,</li><li>    nni.ConvBnReLU1d: nniqat.ConvBnReLU1d,</li><li>    nni.ConvBnReLU2d: nniqat.ConvBnReLU2d,</li><li>    nni.ConvReLU2d: nniqat.ConvReLU2d,</li><li>    nni.LinearReLU: nniqat.LinearReLU</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<p>
or pass in a mapping that includes the new qat module in <a href="https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantize.py#L390" class="inline-onebox" rel="noopener nofollow ugc">pytorch/quantize.py at master · pytorch/pytorch · GitHub</a></p>

username: thyeros
cooked: <p>Hi, Jerry, thanks for sharing that. Yes, I also saw the mapping table, but (I should have been more clear on this): For a custom layer, I need to make a corresponding QAT version of that customer layer. Is there any particular requirement for such QAT versions (in doc or example)? such as a particular set of functions or attributes? One of them looks like ‘from_float’.</p>

username: thyeros
cooked: <p>I figured it out. For anyone with the same interest, hope this template helpful. Jerry, please comment if anything missing here.</p>
<pre><code class="lang-auto">class qatCustom (nn.Module):
    def __init__(..., qconfig=None):
        super().__init__(....)
        self.qconifg = qconfig
        #skip this if there is no learnable parm
        self.weight_fake_quant = qconfig.weight()

    def forward(self, input):
        #implement the forward but with qweight
        qweight = self.weight_fake_quant(self.weight)
        #do compute
        return _forward_imple(input, qweight)

    @classmethod
    def from_float(cls, mod):
       qconfig = mod.qconfig
       qmod = cls(qmod constructor inputs, qconfig=qconfig)    
       qmod.weight = mod.weight
       return qmod
     
torch.quantization.quantization_mappings.register_qat_module_mapping(Custom, qatCustom) 
</code></pre>

username: Wang_Yihan
cooked: <p>What should the qmod constructor inputs be? I defined a threshold comparison function for quantization after training, but encountered a lot of difficulties</p>
<aside class="quote no-group" data-username="thyeros" data-post="4" data-topic="110966" data-full="true">
<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/thyeros/48/12827_2.png" class="avatar"> thyeros:</div>
<blockquote>
<p>I figured it out. For anyone with the same interest, hope this template helpful. Jerry, please comment if anything missing here.</p>
<pre><code class="lang-auto">class qatCustom (nn.Module):
    def __init__(..., qconfig=None):
        super().__init__(....)
        self.qconifg = qconfig
        #skip this if there is no learnable parm
        self.weight_fake_quant = qconfig.weight()

    def forward(self, input):
        #implement the forward but with qweight
        qweight = self.weight_fake_quant(self.weight)
        #do compute
        return _forward_imple(input, qweight)

    @classmethod
    def from_float(cls, mod):
       qconfig = mod.qconfig
       qmod = cls(qmod constructor inputs, qconfig=qconfig)    
       qmod.weight = mod.weight
       return qmod
     
torch.quantization.quantization_mappings.register_qat_module_mapping(Custom, qatCustom) 
</code></pre>
</blockquote>
</aside>
<pre><code class="lang-auto">class Threshold(nn.Module):
    def __init__(self, num_channels, qconfig=None):
        super().__init__()
        self.num_channels = num_channels
        self.threshold_plus = nn.Parameter(torch.ones(size=[num_channels]), requires_grad=False)
        self.threshold_minus = nn.Parameter(torch.ones(size=[num_channels]), requires_grad=False)
        self.threshold_plus_sign = nn.Parameter(torch.ones(size=[num_channels]), requires_grad=False)
        self.threshold_minus_sign = nn.Parameter(torch.ones(size=[num_channels]), requires_grad=False)
        # self.quant = torch.quantization.QuantStub()
        # self.dequant = torch.quantization.DeQuantStub()
        self.qconifg = qconfig
        self.weight_fake_quant = qconfig.weight()
    def forward(self, x):
        # x = self.quant(x)
        # assert x.size(1) == self.num_channels, f"Expected {self.num_channels} channels, but got {x.size(1)}"
        # Apply thresholding to each channel
        output = []
        for i in range(self.num_channels):
            channel = x[:, i, :]
            threshold_p = self.weight_fake_quant(self.threshold_plus[i])
            # threshold_p = self.threshold_plus[i]
            threshold_m = self.threshold_minus[i]
            p_sign = self.threshold_plus_sign[i]
            m_sign = self.threshold_minus_sign[i]
            channel_output = torch.where(channel &gt;= 0, torch.where(channel &gt; threshold_p, p_sign * torch.tensor(1.0),
                           p_sign * torch.tensor(-1.0)),
                            torch.where(channel &gt; threshold_m, m_sign * torch.tensor(1.0),
                             m_sign * torch.tensor(-1.0)))
            output.append(channel_output)
        # Stack the thresholded channels along the channel dimension
        output = torch.stack(output, dim=1)
        # output = self.dequant(output)
        return output
    @classmethod
    def from_float(cls, mod):
        qconfig = mod.qconfig
        qmod = cls(qmod constructor inputs, qconfig = qconfig)
        qmod.weight = mod.weight
        return qmod
</code></pre>