Hussain5/Quantized-Mobilenet

1import torch2import torch.nn as nn3import torch.nn.functional as F4from QuantizedMobileNetBlock import QuantizedMobileNetBlock5 6 7# Build the full model8class QuantizedMobileNet(nn.Module):9    def __init__(self, config):10        super().__init__()11        self.blocks = nn.ModuleList()12        input_channels = 313        for out_channels, stride, bits in config:14            block = QuantizedMobileNetBlock(input_channels, out_channels, stride, bits)15            self.blocks.append(block)16            input_channels = out_channels17        self.classifier = nn.Linear(input_channels, 10)18 19    def forward(self, x):20        for block in self.blocks:21            x = block(x)22        x = F.adaptive_avg_pool2d(x, 1)23        x = torch.flatten(x, 1)24        x = self.classifier(x)25        return x26 27    def total_bitops(self, input_size):28        total = 029        current_size = input_size30        for block in self.blocks:31            ops, current_size = block.bitops(current_size)32            total += ops33        return total34 35    def bitops_per_layer(self, input_size):36        layerwise = []37        current_size = input_size38        for idx, block in enumerate(self.blocks):39            ops, current_size = block.bitops(current_size)40            layerwise.append((idx + 1, ops))41        return layerwise