Setting number of bits for Weights and activations

username: Sami_Hassan
cooked: <p>I would like to change the number of bits for weights and activation during testing. The network is pretrained and .pth is available which is usually 32 bit as per my understanding. Lets say i have the following network which is already trained.</p>
<pre><code class="lang-auto">class AlexNet(nn.Module):

    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


model = AlexNet()
state_dict = load_state_dict_from_url(model_urls['alexnet'],
                                            progress=True)
model.load_state_dict(state_dict)

</code></pre>
<p>For Testing:</p>
<p>I want to change the activation’s of this network to be 9-8-5-5-7 bits which are the  Relu’s in self.features  per my understanding and weights to be 11 bit in all convolution layers.</p>
<p>I also want to change to Fully Connected layers weigh to be 10-9-9 bit respectively.</p>
<p>Can someone please point me to the right direction or share a code snippet. Thanking in anticipation.</p>

username: albanD
cooked: <p>Hi,</p>
<p>What are these numbers of bits? Floating points are usually 64, 32 or 16 bits. Do you refer to other encodings here?</p>

username: Sami_Hassan
cooked: <p>Hello Thanks for your reply. They refer to quantized bits as per my understanding. <a href="https://arxiv.org/pdf/1706.07853.pdf" rel="nofollow noopener">[page-4]</a></p>
<p>Currently i have tried this solution <a href="https://github.com/aaron-xichen/pytorch-playground" rel="nofollow noopener">[inside utee folder quant file]</a></p>
<pre><code class="lang-auto">arr = [11,11,11,11,11,11,11,11,11,11,10,10,9,9,9,9]
count = 0
bit_num = OrderedDict()

for k, v in state_dict.items():
    bit_num[k] = arr[count]    
    count += 1

state_dict_quant = OrderedDict()

#state_dict_quant['features.0.weight'] = quant.log_minmax_quantize(model.features[0].weight, bits=11)

for k, v in state_dict.items():
    print(k,v.type())
    state_dict_quant[k] = quant.log_minmax_quantize(state_dict[k], bits=bit_num[k])
</code></pre>
<p>Looks like it is doing the job, though i still have to think how to modify bits for activation, as these are calculated on run time ? . May be “hook” can help.</p>

username: albanD
cooked: <p>You can take a look at the <a href="https://pytorch.org/docs/stable/quantization.html">quantization</a> module to see if it fits your needs?</p>

username: Sami_Hassan
cooked: <p>I am not sure how can this module help in setting custom number of bits, is it possible using quantization module ?.</p>
<p>Also can you please help me with a code snippet to change number of bits for the activation ?  Currently i am using this one but it doesn’t seem to work.</p>
<pre><code class="lang-auto">arr1 = [9,8,5,5,7]     # activations quantized 
#arr1 = [2,2,2,2,2]

def quantizee(self,input,output):
    #print("Hello")
    #print('Inside ' + self.__class__.__name__ + ' forward')
    #print('input: ', type(input))
    #print('output: ', type(output))
    
    print('output size:', output.data.size())
    print(arr1[quantizee.counter])
    #out = output.detach()
    #print('out: ', type(out))
    output = quant.min_max_quantize(output.data,arr1[quantizee.counter])
    quantizee.counter += 1
    if (quantizee.counter == 5):
        quantizee.counter = 0



quantizee.counter = 0

count1 = 0          # number of relus you want to quantize count 1 is that number
for name, module in model.named_modules():
    if (isinstance (module, torch.nn.modules.Conv2d) &amp; (count1 &lt; 5)):
        module.register_forward_hook(quantizee)
        count1 += 1'''
</code></pre>
<p>i have tried this one too but no success.</p>
<pre><code class="lang-auto">#arr1 = [9,8,5,5,7]
arr1 = [2,2,2,2,2]
count1 = 0
l = OrderedDict()
for k, v in state_dict_quant.items():
    if isinstance(v, (nn.ReLU)):
        quant_layer = quant.NormalQuant('{}_quant'.format(k), bits=arr1[count1], quant_func=log_minmax_quantize)
        l['{}_{}_quant'.format(k, type)] = quant_layer
        count1 += 1
        if (count1 == 5):
            count1= 0
    
    else:
        l[k]= v
        
</code></pre>