How does PyTorch implement Quantization?

username: Eta_C
cooked: <p>Hello, I want to know how does PyTorch implement quantization without rewrite the layer definition.<br>
Actually, because of certain reasons, I have to write a quantization framework with PyTorch for my project. Here is an simple example of my implementaion:</p>
<pre><code class="lang-python">def quantization_decorator(func):
    @wraps(func)
    def wrapper(_, x):
        do_quantization(x)
        x = func(x)
        do_quantization(x)
        return x
    return wrapper

class Conv2d(nn.Conv2d):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    
    @quantization_decorator
    def forward(self, x):
        super().forward(x)
</code></pre>
<p>Since my quantization method is defined as a decorator,<br>
<strong>ADWANTAGE</strong>: I can add it to <code>forward</code> to choose whether to quantize this module.<br>
<strong>DISAWANTAGE</strong>: I have to rewrite each layer, e.g. <code>class Conv2d</code> in my example.</p>
<p>I want to know how does PyTorch implement quantization, so that I can optimize my framework (simpler, faster and more customizable).</p>
<p>Thank you.</p>

username: Eta_C
cooked: <ul>
<li>My quantization method is different from PyTorch.</li>
<li>My custom layers have to be supported.</li>
</ul>
<p>That’s why I have to write a new implementation.<img src="https://discuss.pytorch.org/images/emoji/apple/see_no_evil.png?v=9" title=":see_no_evil:" class="emoji" alt=":see_no_evil:"></p>

username: dskhudia
cooked: <p>Hi <a class="mention" href="/u/eta_c">@Eta_C</a>,</p>
<p>Please look at the flow of operation for quantization here: <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html#post-training-static-quantization" rel="nofollow noopener">https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html#post-training-static-quantization</a></p>
<p>The main steps for post training quantization are: 1) fusing of modules (e.g., conv; bn; relu =&gt; conv_bn_relu) 2) Observing tensor values to quantize tensors 3) Actual replacing of modules from float to quantized.</p>

username: Eta_C
cooked: <p>I know how to use PyTorch quantization tools.<br>
I want to know how does PyTorch implement it.<br>
OK, maybe I have to read some source code about it…</p>

username: dskhudia
cooked: <p>Yes. Reading source code on how it’s implemented is a good way. Here are some pointers for the code: Python related: <a href="https://github.com/pytorch/pytorch/tree/master/torch/quantization" rel="nofollow noopener">https://github.com/pytorch/pytorch/tree/master/torch/quantization</a><br>
C++ kernels: <a href="https://github.com/pytorch/pytorch/tree/master/aten/src/ATen/native/quantized/cpu" rel="nofollow noopener">https://github.com/pytorch/pytorch/tree/master/aten/src/ATen/native/quantized/cpu</a></p>
<p>You may also want to check the recent changes to these file and comments on related PRs to get some context.</p>

username: Eta_C
cooked: <p>Thank you.<br>
I will take the time to read the source code carefully.<img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=9" title=":grinning:" class="emoji" alt=":grinning:"></p>

username: remya
cooked: <p>I was trying to quantize a vgg model , following the steps given in <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html" rel="nofollow noopener">https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html</a></p>
<p>I did the first step as follows:</p>
<p>import torch.nn as nn<br>
<span class="hashtag">#from</span> .utils import load_state_dict_from_url</p>
<p><strong>all</strong> = [<br>
‘Cifar_VGG’, ‘cifar_rvgg11’, ‘cifar_rvgg11_bn’, ‘cifar_vgg11’, ‘cifar_vgg11_bn’, ‘cifar_vgg13’, ‘cifar_vgg13_bn’, ‘cifar_vgg16’, ‘cifar_vgg16_bn’,<br>
‘cifar_vgg19_bn’, ‘cifar_vgg19’,<br>
]</p>
<p>model_urls = {</p>
<pre><code>'cifar_vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
</code></pre>
<p>}</p>
<p>class Cifar_VGG(nn.Module):</p>
<pre><code>def __init__(self, features, num_classes=1000, init_weights=True):
    super(Cifar_VGG, self).__init__()
    self.features = features
    self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
    self.quant = QuantStub()
    self.dequant = DeQuantStub()
    self.classifier = nn.Sequential(
        nn.Linear(512, 512),
        nn.BatchNorm1d(512),
        nn.ReLU(True),
        nn.Linear(512, num_classes),
    )
    if init_weights:
        self._initialize_weights()

def forward(self, x):
    x = self.quant(x)
    x = self.features(x)
    x = self.avgpool(x)
    x = x.view(x.size(0), -1)
    x = self.classifier(x)
    x = self.dequant(x)

    return x

def _initialize_weights(self):
    for m in self.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, 0, 0.01)
            nn.init.constant_(m.bias, 0)

def fuse_model(self):
    torch.quantization.fuse_modules(self, [['conv2d', 'BatchNorm2d', 'ReLU'],
                                        ],inplace=True)
</code></pre>
<p>def make_layers(cfg, batch_norm=False):<br>
layers = []<br>
in_channels = 3<br>
for v in cfg:<br>
if v == ‘M’:<br>
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]<br>
else:<br>
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)<br>
if batch_norm:<br>
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]<br>
else:<br>
layers += [conv2d, nn.ReLU(inplace=True)]<br>
in_channels = v<br>
return nn.Sequential(*layers)</p>
<p>cfgs = {<br>
‘RA’:[512,‘M’, 512, ‘M’, 512, 512, ‘M’, 512, 512, ‘M’, 512, 512, ‘M’],<br>
‘A’: [64, ‘M’, 128, ‘M’, 256, 256, ‘M’, 512, 512, ‘M’, 512, 512, ‘M’],<br>
‘B’: [64, 64, ‘M’, 128, 128, ‘M’, 256, 256, ‘M’, 512, 512, ‘M’, 512, 512, ‘M’],<br>
‘D’: [64, 64, ‘M’, 128, 128, ‘M’, 256, 256, 256, ‘M’, 512, 512, 512, ‘M’, 512, 512, 512, ‘M’],<br>
‘E’: [64, 64, ‘M’, 128, 128, ‘M’, 256, 256, 256, 256, ‘M’, 512, 512, 512, 512, ‘M’, 512, 512, 512, 512, ‘M’],<br>
}</p>
<p>def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs):<br>
if pretrained:<br>
kwargs[‘init_weights’] = False<br>
model = Cifar_VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)<br>
if pretrained:<br>
state_dict = load_state_dict_from_url(model_urls[arch],<br>
progress=progress)<br>
model.load_state_dict(state_dict)<br>
return model</p>
<p>def cifar_rvgg11_bn(pretrained=False, progress=True, **kwargs):<br>
return _vgg(‘cifar_rvgg11’, ‘RA’, True, pretrained, progress, **kwargs)</p>
<p>net.fuse_model()<br>
Gives the following error</p>
<p>‘Cifar_VGG’ object has no attribute ‘conv2d’</p>

username: jerryzh168
cooked: <p>I think the fusion is not defined correctly for your model, please read the fusion section of the tutorial again and see if you can find the problem</p>

username: remya
cooked: <p>Now I am getting the following error<br>
RuntimeError: No function is registered for schema aten::native_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps) -&gt; (Tensor, Tensor, Tensor) on tensor type QuantizedCPUTensorId; available functions are CPUTensorId, CUDATensorId, MkldnnCPUTensorId, VariableTensorId</p>

username: remya
cooked: <p>How to merge batch_norm layer with conv layer?</p>

username: remya
cooked: <p>Now i am trying ResNet18<br>
def fuse_model(self):<br>
for m in self.modules():<br>
if type(m) == BasicBlock:<br>
torch.quantization.fuse_modules(m, [[‘conv1’, ‘bn1’, ‘relu’],[‘conv2’, ‘bn2’]], inplace=True)<br>
conv = getattr(m, ‘downsample’)<br>
if(conv) :<br>
torch.quantization.fuse_modules(conv, [‘0’, ‘1’], inplace=True)</p>
<p>But fusion is not giving any error.<br>
But the infernce shows the error (which i posted above)</p>
<p>Got it !!!<br>
I missed the following:<br>
torch.quantization.fuse_modules(self, [‘conv1’, ‘bn1’, ‘relu’], inplace=True)</p>
<p>Thank you</p>