BitNetMCU.py

import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F

# @cpldcpu 2024-June-2

class Activation(nn.Module):
    def __init__(self, mode='hardswish'):
        super(Activation, self).__init__()
        self.mode = mode

    def forward(self, x):
        if self.mode == 'hardswish':
            return x * F.relu6(x + 3) / 6
        elif self.mode == 'GeLU':
            return F.gelu(x)
        elif self.mode == 'ReLU':
            return F.relu(x)
        elif self.mode == 'ReLU2':
            return x * F.relu(x)
        elif self.mode == 'ReLU2swish':
            return x * F.relu(x + 3)
        else:
            raise ValueError(f"Unknown activation mode: {self.mode}")

class BitQuant:
    """
    Class to handle quantization of activations and weights.

    Quantization Types:
    - Binary         : 1 bit
    - Ternary        : 1.58 bits
    - BinaryBalanced : 1 bit, weights are balanced around zero
    - 2bitsym        : 2 bit symmetric
    - 4bitsym        : 4 bit symmetric
    - FP130          : 4 bit shift encoding
    - NF4            : 4 bit non-linear quantization (NormalFloat4)
    - 8bit           : 8 bit

    WScale
    - PerTensor      : The weight scaling is calculated per Tensor
    - PerOutput      : The weight scaling is calculated per Output

    Implementation was initially based on:
    https://github.com/microsoft/unilm/blob/master/bitnet/The-Era-of-1-bit-LLMs__Training_Tips_Code_FAQ.pdf

    """
    def __init__(self, QuantType='Binary', WScale='PerTensor'):
        self.QuantType = QuantType
        self.WScale = WScale
        self.s = torch.nn.Parameter(torch.tensor(1.0))
        self.s.requires_grad = False  # no gradient for clipping scalar

        # Activation quantization (set by the model builder / training script):
        #   act_bits : 8 (default, per-token max, signed -128..127) or 4 (16 linear levels 0..15 for nonnegative tokens,
        #              -8..7 for signed tokens, per-token max)
        #   act_pow2 : True -> the per-token scale is rounded down to a power of two and codes are truncated (floor),
        #              which is exactly what the MCU ShiftNorm does (x >> shift). False -> exact max scaling + rounding.
        #   table_scale : >0 -> nonuniform (NF4) weight levels are rounded to multiples of 1/table_scale, i.e. to the
        #              integer grid of the int16 product table T[a,w] = a * round(level_w * table_scale) on the MCU.
        self.act_bits = 8
        self.act_pow2 = False
        self.act_unsigned = False   # True for layers fed by a ReLU (set by the model builder): nonnegative token codes
        self.act_pow2_round = False # pow2 mode: round-to-nearest ((acc + 2^(sh-1)) >> sh on the device) instead of truncation
        # pow2 mode removes every non-power-of-two scale from the graph, so the network cannot set its own softmax
        # temperature any more (hidden layers are scale-free through ReLU + ShiftNorm, but the logits are not).
        # The output layer therefore multiplies its logits by a learnable temperature exp(log_logit_scale); this is
        # a positive per-model constant and invisible to the argmax on the device.
        self.is_output = False
        # act_group > 1: the per-token max (and hence the shift) is shared by groups of `act_group` consecutive tokens
        # (e.g. the four 8x8 blocks of one image in PatchMNIST). Keeps the relative magnitude of the blocks intact
        # through a power-of-two shift; on the device this is one max over the group's activations.
        self.act_group = 1
        # Requantization refinement for the pow2 (MCU-exact) path; integer-exact on the device:
        #   act_mantissa : 'none' (plain shift) | 'max' (per token m in 8..15 so that m*max/8 just fits the range,
        #                  one 4x4-bit multiply + >>3 per activation) | 'recipN' (per-token N-bit reciprocal
        #                  s = floor(qmax*2^t/max), code = (acc*s) >> t; N = 8 matches the exact max scaling and is
        #                  the recommended device mode: one N-bit multiply per activation + one reciprocal per token).
        self.act_mantissa = 'none'
        self.log_logit_scale = torch.nn.Parameter(torch.tensor(0.0))   # init 1.0, learned by Adam
        self.table_scale = 0

        if self.QuantType in ['Binary', 'BinarySym']:
            self.bpw = 1
        elif self.QuantType in ['2bitsym', 'NF2']:
            self.bpw = 2
        elif self.QuantType in ['Ternary']:
            self.bpw = 1.6
        elif self.QuantType in ['4bit', '4bitsym', 'sint4', 'FP130' , 'NF4']:
            self.bpw = 4
        elif self.QuantType == '5bitsym':
            self.bpw = 5
        elif self.QuantType == '8bit':
            self.bpw = 8
        else:
            raise AssertionError(f"Invalid QuantType: {self.QuantType}")

        if not self.WScale in ['PerOutput', 'PerTensor']:
            raise AssertionError(f"Invalid WScale: {self.WScale}. Expected one of: 'PerTensor', 'PerOutput'")

    NF4_LEVELS = [-1.0, -0.6962, -0.5251, -0.3949, -0.2844, -0.1848, -0.0911, 0.0,
                  0.0796, 0.1609, 0.2461, 0.3379, 0.4407, 0.5626, 0.723, 1.0]
    # 4-level Lloyd-Max quantizer for a Gaussian (+-0.4528 sigma, +-1.510 sigma), normalized to the outer level:
    NF2_LEVELS = [-1.0, -0.2998, 0.2998, 1.0]

    def code_levels(self, device=None):
        """Codebook levels in units of the quantization scale, i.e. u = w*scale takes values in this set.
        For NF4 the levels are optionally rounded to the int16 product-table grid (table_scale)."""
        if self.QuantType in ('NF4', 'NF2'):
            lv = torch.tensor(self.NF4_LEVELS if self.QuantType == 'NF4' else self.NF2_LEVELS, device=device)
            if self.table_scale > 0:
                lv = torch.round(lv * self.table_scale) / self.table_scale
            return lv
        elif self.QuantType in ['2bitsym', '4bitsym', '5bitsym']:
            n = 2 ** (self.bpw - 1)
            return torch.arange(-n, n, device=device).float() + 0.5
        elif self.QuantType in ['4bit', 'sint4', '8bit']:
            n = 2 ** (self.bpw - 1)
            return torch.arange(-n, n, device=device).float()
        else:
            raise AssertionError(f"code_levels not defined for QuantType {self.QuantType}")

    def octav_cb(self, tensor, num_iterations=10, s=-1):
        """Codebook-aware OCTAV: same fixed-point iteration as octav(), but the per-unit-s^2 MSE of the in-range
        weights is measured empirically on the actual level set instead of the uniform-grid constant 4^-b/3.
        For a uniform grid this recovers ~4^-b/3; for NF4 etc. it gives the matched clip."""
        levels = self.code_levels(tensor.device)
        bound = levels.abs().max()          # outermost level in units of scale; s = bound / scale
        if s < 0:
            s = tensor.abs().mean().clamp_(min=1e-5) * 0.25
        a = tensor.abs()
        for _ in range(num_iterations):
            le = a <= s
            gt = ~le
            u = tensor[le] * (bound / s)     # in-range weights in level units
            q, _ = self.quantize_list(u, levels)
            C = ((u - q) ** 2).mean() / bound ** 2 if u.numel() > 0 else torch.tensor(4.0 ** -self.bpw / 3, device=tensor.device)
            numerator = torch.sum(a[gt])
            denominator = C * le.sum() + gt.sum()
            s = numerator / denominator.clamp(min=1)
        return s

    # Octave optimum clipping algorithm (C. Sakr et al., 2022)
    # see https://arxiv.org/abs/2206.06501
    def octav(self, tensor, num_iterations=10, s=-1):
        if s<0:
            # s = torch.sum(torch.abs(tensor)) / torch.sum(tensor != 0)
            s = tensor.abs().mean().clamp_(min=1e-5) * 0.25 # blind estimate as starting point
        for _ in range(num_iterations):
            indicator_le = (torch.abs(tensor) <= s).float()
            indicator_gt = (torch.abs(tensor) > s).float()
            numerator = torch.sum(torch.abs(tensor) * indicator_gt)
            denominator = (4**-self.bpw / 3) * torch.sum(indicator_le) + torch.sum(indicator_gt)
            s = numerator / denominator
        return s

    def update_clipping_scalar(self, w, algorithm='octav', quantscale=0.25):
        """
        Update the weight scale factor for the quantization.
        Args:
            w:          a weight tensor with shape [d, k]
            algorithm:  clipping algorithm to use
                'octav' : Octave optimum clipping algorithm
                'prop'  : Proportional clipping algorithm
        Returns:
            s:          updated clipping scalar for the quantization
        """

        s= self.s

        if algorithm == 'octav':
            if self.WScale=='PerOutput':
                s = torch.stack([self.octav(row, 10) for row in w])
            else:
                s = self.octav(w, 10, s)
        elif algorithm == 'octav_cb':
            if self.WScale=='PerOutput':
                s = torch.stack([self.octav_cb(row, 10) for row in w])
            else:
                s = self.octav_cb(w, 10, s)
        elif algorithm == 'prop':
            if self.WScale=='PerOutput':
                s = w.abs().max(dim=-1, keepdim=True)[0].clamp_(min=1e-5) / quantscale
            else:
                s = w.abs().mean().clamp_(min=1e-5) / quantscale
        else:
            raise AssertionError(f"Invalid algorithm: {algorithm}. Expected one of: 'octav', 'octav_cb', 'prop'")

        self.s = torch.nn.Parameter(s)
        self.s.requires_grad = False  # no gradient for clipping scalar

        return s

    def activation_quant(self, x):
        """ Per-token activation quantization (no grouping needed).
        Args:
        x: an activation tensor with shape [n, d] (last dim = token)
        Returns:
        y:     integer activation codes (as float tensor)
        scale: per-token scale factor, x_quant = y / scale
        """
        xmax = x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
        if self.act_group > 1:
            g = self.act_group
            xmax = xmax.reshape(-1, g, 1).max(dim=1, keepdim=True).values.expand(-1, g, 1).reshape(-1, 1)
        if self.act_bits == 8:
            qmin, qmax = -128, 127
        elif self.act_bits == 4:
            qmin, qmax = (0, 15) if self.act_unsigned else (-8, 7)
        else:
            raise AssertionError(f"Unsupported act_bits: {self.act_bits}")

        if self.act_pow2 and self.act_bits in (4, 8) and str(self.act_mantissa).startswith('recip'):
            # Per-token fixed-point reciprocal: s = floor(qmax * 2^t / max) with s in [2^(b-1), 2^b) (b = bits of s),
            # code = (acc * s) >> t  (or rounded).  One reciprocal per token, one b-bit multiply per activation; approaches
            # the exact per-token max scaling as b grows.  Computed in float64: x is integer * 2^k, acc*s < 2^40.
            b = int(self.act_mantissa[5:])
            xd = x.double(); xm = xmax.double()
            t = torch.ceil(torch.log2(xm / qmax)) + (b - 1)          # 2^t*qmax/max in [2^(b-1), 2^b]
            p2t = torch.exp2(t)
            sc = torch.floor(qmax * p2t / xm).clamp_(max=2 ** b - 1)   # integer s (the exact-power-of-two case is clamped)
            prod = xd * sc / p2t
            y = torch.floor(prod + 0.5) if self.act_pow2_round else torch.floor(prod)
            return y.clamp_(qmin, qmax).float(), (sc / p2t).float()

        if self.act_pow2:
            # MCU ShiftNorm: power-of-two scale bringing the token max into [(qmax+1)/2, qmax+1), then truncation.
            # (x is integer-valued times a power of two in this mode, so floor(x*scale) == acc >> shift exactly.)
            top = qmax + 1
            # largest k with xmax * 2^k < top  (i.e. xmax*2^k in [top/2, top))
            scale = torch.exp2(torch.ceil(torch.log2(top / xmax)) - 1)
            if self.act_pow2_round:
                xs = torch.floor(x * scale + 0.5)
            else:
                xs = torch.floor(x * scale)
            xs = xs.clamp_(qmin, top - 1)
            gain = scale
            if self.act_bits == 4 and self.act_mantissa == 'max':
                # per token: m = floor((8*top - 1) / max_s) clamped to 8..15, so that m*max_s/8 < top
                smax = xs.max(dim=-1, keepdim=True).values.clamp_(min=1)
                if self.act_group > 1:
                    g = self.act_group
                    smax = smax.reshape(-1, g, 1).max(dim=1, keepdim=True).values.expand(-1, g, 1).reshape(-1, 1)
                m = torch.floor((8 * top - 1) / smax).clamp_(8, 15)
                xs = torch.floor(xs * m / 8)              # (x_s * m) >> 3, exact (integers < 2^24)
                gain = gain * m / 8
            y = xs.clamp(qmin, qmax)
            return y, gain

        else:
            scale = qmax / xmax
            y = (x * scale).round().clamp_(qmin, qmax)
        return y, scale

    def quant_forward_pow2(self, x, w):
        """Exact-MCU forward used when act_pow2 is set: no RMS normalization (ShiftNorm is the normalization),
        activations dequantized by the power of two 2^act_bits (16 or 128) and weights kept in level units
        (u / bound, bound = 2^(bpw-1) or 1 for NF4), so every value in the graph is an integer times a power of two
        and the float computation reproduces the integer kernel bit-for-bit (as long as |int| < 2^24).
        Returns (x_quant, w_quant) with straight-through gradients."""
        x_int, x_scale = self.activation_quant(x)
        qnorm = 2.0 ** self.act_bits / 2 if not self.act_unsigned else 2.0 ** self.act_bits
        # signed: codes/128 (8-bit) or /8 (4-bit) ; unsigned: codes/16
        x_quant = x * (x_scale / qnorm) + (x_int / qnorm - x * (x_scale / qnorm)).detach()
        w_int, w_scale, _ = self.weight_quant(w)
        bound = 1.0 if self.QuantType in ('NF4', 'NF2') else 2.0 ** (self.bpw - 1)
        w_quant = w * (w_scale / bound) + (w_int / bound - w * (w_scale / bound)).detach()
        return x_quant, w_quant

    def weight_quant(self, w):
        """ Per-tensor quantization.
        Args:
        w: a weight tensor with shape [d, k]
        Returns:
        u:     a quantized weight with shape [d, k]
        scale: scale factor for the quantization
        bpw:   bit per weight
        """

        if self.QuantType == 'FP130':
           scale = 128.0 / self.s
        elif self.QuantType in ('NF4', 'NF2'):
            scale = 1.0 / self.s
        elif self.QuantType == 'Ternary': # 1.58bits
            # scale = 1.0 / self.s
            scale = 1.0 / w.abs().mean().clamp_(min=1e-5)
        else:
            scale = (2.0**(self.bpw-1)) / self.s

        if self.QuantType == 'Ternary': # 1.58bits
            u = (w * scale ).round().clamp_(-1, 1)
        elif self.QuantType == 'Binary': # 1 bit
            e = w.mean()
            u = (w - e).sign()
        elif self.QuantType == 'BinarySym': # 1 bit
            u = w.sign()
        elif self.QuantType == '2bitsym':
            u = ((w * scale - 0.5).round().clamp_(-2, 1) + 0.5)
        elif self.QuantType == '4bit': # 4 bit in one-complement encoding for inference with multiplication
            # u = (w * scale).round().clamp_(-8, 7) # no convergence with this?!
            u = ((w * scale - 0.01).round().clamp_(-8, 7) + 0.01)
        elif self.QuantType == '4bitsym':
            u = ((w * scale - 0.5).round().clamp_(-8, 7) + 0.5)
        elif self.QuantType == 'sint4':  # plain signed int4 levels -8..7 (with zero) for multiply-based kernels (RV32EMC)
            u = (w * scale).round().clamp_(-8, 7)
        elif self.QuantType ==  'FP130': # encoding (F1.3.0) : S * ( 2^E3 + 1) -> min 2^0 = 1, max 2^7 = 128
            e = ((w * scale).abs()).log2().floor().clamp_(0, 7)
            u = w.sign()*(e.exp2())
        elif self.QuantType in ('NF4', 'NF2'):
            # normal-float levels (16 for 4 bits, 4 for 2 bits)
            levels = self.code_levels(w.device)
            u , _ = self.quantize_list(w * scale, levels)
        elif self.QuantType == '5bitsym':
            u = ((w * scale - 0.5).round().clamp_(-16, 15) + 0.5)
        elif self.QuantType == '8bit': # -128 to 127
            u = (w * scale).round().clamp_(-128, 127)
        else:
            raise AssertionError(f"Invalid QuantType: {self.QuantType}. Expected one of: 'Binary', 'BinaryBalanced', '2bitsym', '4bitsym', '8bit'")

        return u, scale, self.bpw

    def quantize_list(self, x, levels):
        """
        Quantize the input tensor x to the nearest level in the levels list.
        """
        # Compute the absolute difference between x and each level
        diff = torch.abs(x.unsqueeze(-1) - levels)
        # Find the index of the closest level for each element in x
        indices = torch.argmin(diff, dim=-1)

        return levels[indices], indices

class BitLinear(nn.Linear, BitQuant):
    """
    Linear fully connected layer with quantization aware training and normalization.
    Configurable quantization and normalization types.

    Normalization Types:
    - RMS            : Root Mean Square
    - Lin            : L1 Norm
    - BatchNorm      : Batch Normalization
    - LayerNorm      : Layer Normalization

    This is not optimized for speed or efficiency...

    @cpldcpu 2024-March-24
    """
    def __init__(self, in_features, out_features, bias=False, QuantType='Binary', WScale='PerTensor', NormType='RMS'):
        nn.Linear.__init__(self, in_features, out_features, bias=False)
        BitQuant.__init__(self, QuantType, WScale)

        self.NormType = NormType

    def forward(self, x):
        """
        Args:
        x: an input tensor with shape [n, d]
        Returns:
        y: an output tensor with shape [n, k]
        """
        w = self.weight # a weight tensor with shape [d, k]

        if self.act_pow2 and self.QuantType != 'None':
            x_quant, w_quant = self.quant_forward_pow2(x, w)
            y = F.linear(x_quant, w_quant)
            if self.is_output:
                y = y * self.log_logit_scale.exp()
            return y

        x_norm = self.Normalize(x)

        if self.QuantType == 'None':
            y = F.linear(x_norm, w)
        else:
            # A trick for implementing Straight-Through-Estimator (STE) using detach()
            x_int, x_scale = self.activation_quant(x_norm)
            x_quant = x_norm + ( x_int / x_scale - x_norm).detach()

            w_int, w_scale, _ = self.weight_quant(w)
            w_quant = w + (w_int / w_scale - w).detach()

            y = F.linear(x_quant, w_quant)
        return y

    def Normalize(self, x):
        """ Normalization. Normalizes along the last dimension -> different normalization value for each activation vector.
        Args:
        x: an input tensor with shape [n, d]
        Returns:
        y: a normalized tensor with shape [n, d]
        """
        if self.NormType == 'RMS':
            y = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True)).clamp(min=1e-8)  # eps guard: all-zero tokens
            z =x / y
        elif self.NormType == 'Lin':
            y = torch.mean(torch.abs(x), dim=-1, keepdim=True)
            z =x / y
        elif self.NormType == 'BatchNorm':
            # BatchNorm: normalize across batch dimension
            batch_mean = torch.mean(x, dim=0, keepdim=True)
            batch_var = torch.var(x, dim=0, keepdim=True, unbiased=False)
            z = (x - batch_mean) / torch.sqrt(batch_var + 1e-5)
        elif self.NormType == 'LayerNorm':
            # LayerNorm: normalize across feature dimension
            layer_mean = torch.mean(x, dim=-1, keepdim=True)
            layer_var = torch.var(x, dim=-1, keepdim=True, unbiased=False)
            z = (x - layer_mean) / torch.sqrt(layer_var + 1e-5)
        else:
            raise AssertionError(f"Invalid NormType: {self.NormType}. Expected one of: 'RMS', 'Lin', 'BatchNorm', 'LayerNorm'")
        return z

class BitConv2d(nn.Conv2d, BitQuant):
    """
    2D convolution layer with quantization aware training and normalization.
    Configurable quantization and normalization types.

    Normalization Types:
    - RMS            : Root Mean Square
    - None           : No normalization

    @cpldcpu 2024-June-2
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, groups=1,  QuantType='4bitsym', WScale='PerTensor', NormType='RMS'):
        nn.Conv2d.__init__(self,in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False)
        BitQuant.__init__(self, QuantType, WScale)

        self.NormType = NormType
        self.groups = groups
        self.stride = stride
        self.padding = padding

    def forward(self, x):
        """
        Args:
        x: an input tensor with shape [n, d]
        Returns:
        y: an output tensor with shape [n, k]
        """
        w = self.weight # a weight tensor with shape [d, k]

        if self.act_pow2 and self.QuantType != 'None':
            # NOTE: this conv path is implemented for symmetry but was not exercised by a trained model in the
            # accompanying study (all MCU-exact runs used BitLinear); validate before relying on it.
            # per-image token for the conv (max over C,H,W), matching Normalize's reduction
            n = x.shape[0]
            x_quant, w_quant = self.quant_forward_pow2(x.reshape(n, -1), w)
            return F.conv2d(x_quant.reshape_as(x), w_quant, groups=self.groups, stride=self.stride, padding=self.padding, bias=None)

        x_norm = self.Normalize(x)

        if self.QuantType == 'None':
            y = F.conv2d(x_norm, w,  stride=self.stride, padding=self.padding, groups=self.groups )
        else:
            x_int, x_scale = self.activation_quant(x_norm)
            x_quant = x_norm + (x_int / x_scale - x_norm).detach()

            w_int, w_scale, _ = self.weight_quant(w)
            w_quant = w + (w_int / w_scale - w).detach()

            y = F.conv2d(x_quant, w_quant, groups=self.groups, stride=self.stride, padding=self.padding, bias=None)
        return y

    def Normalize(self, x):
            """ Normalization. Normalizes along the last dimension -> different normalization value for each activation vector.
            Args:
            x: an input tensor with shape [n, d]
            Returns:
            y: a normalized tensor with shape [n, d]
            """
            if self.NormType == 'RMS':
                y = torch.sqrt(torch.mean(x**2, dim=(-2,-1), keepdim=True)).clamp(min=1e-8)
                z = x / y
            elif self.NormType == 'None':
                z = x
            else:
                raise AssertionError(f"Invalid NormType: {self.NormType}. Expected one of: 'RMS', 'None'")
            return z

class QuantizedModel:
    """
    This class represents a quantized model. It provides functionality to quantize a given model.
    """

    def __init__(self, model = None):
        self.quantized_model=None
        self.total_bits=0

        if model is not None:
            self.quantized_model, _ = self.quantize(model)

    def totalbits(self):
        """
        Returns the total number of bits used by the quantized model.
        """
        return self.total_bits

    def quantize(self,model):
        """
        This method quantizes the weights of the given model.

        Parameters:
        model (torch.nn.Module): The PyTorch model to be quantized.
        Only the weights of the BitLinear layers are quantized.

        Returns:
        list: A list of dictionaries containing information about each layer of the quantized model.
        int: The total number of bits used by the quantized model.
        """
        quantized_model = []
        totalbits = 0
        for i, layer in enumerate(model.modules()):
            print(i, layer.__class__.__name__)
            if isinstance(layer, BitLinear):
                w         = layer.weight.data

                # print(f'layer: {layer} s:{layer.s})')
                u, scale, bpw = layer.weight_quant(w)
                numscale = 0 # TODO: store scale value for "PerOutput" scaling
                # print (scale)

                totalbits += bpw * u.numel() + numscale * 8
                quantized_weight = u.cpu().numpy()

                layer_info = {
                    'layer_type': 'BitLinear',
                    'layer_order': i,
                    'incoming_weights': quantized_weight.shape[1],
                    'outgoing_weights': quantized_weight.shape[0],
                    'quantized_weights': quantized_weight.tolist(),
                    'WScale': layer.WScale,
                    # 'quantized_scale': scalequant.cpu().numpy().tolist() if layer.WScale=='PerOutput' else [], # TODO: "PerOutput" scaling
                    'bpw': bpw, # bits per weight
                    'quantization_type': layer.QuantType
                }
                quantized_model.append(layer_info)

            elif isinstance(layer, BitConv2d):
                w = layer.weight.data
                u, scale, bpw = layer.weight_quant(w)
                quantized_weight = u.cpu().numpy()

                totalbits += bpw * u.numel()
                layer_info = {
                    'layer_type': 'BitConv2d',
                    'layer_order': i,
                    'in_channels': layer.in_channels,
                    'out_channels': layer.out_channels,
                    'incoming_x': 0,  # will be updated during inference
                    'incoming_y': 0,
                    'outgoing_x': 0,
                    'outgoing_y': 0,
                    'kernel_size': layer.kernel_size,
                    'stride': layer.stride,
                    'padding': layer.padding,
                    'groups': layer.groups,
                    'quantized_weights': quantized_weight.tolist(),
                    'bpw': bpw,
                    'quantization_type': layer.QuantType
                }
                quantized_model.append(layer_info)

            elif isinstance(layer, nn.MaxPool2d):
                layer_info = {
                    'layer_type': 'MaxPool2d',
                    'layer_order': i,
                    'kernel_size': layer.kernel_size,
                    'stride': layer.stride
                }
                quantized_model.append(layer_info)

        self.total_bits = totalbits                
        self.quantized_model = quantized_model

        return quantized_model, totalbits

    def inference_quantized(self, input_data):
        """
        This function performs inference on the given quantized model with the provided input data.

        Parameters:
        quantized_model (list): A list of dictionaries containing information about each layer of the quantized model.
        input_data (torch.Tensor): The input data to be used for inference.

        Returns:
        torch.Tensor: The output of the model after performing inference.
        """

        if not self.quantized_model:
            raise ValueError("quantized_model is empty or None")

        scale = 127.0 / np.maximum(np.abs(input_data).max(axis=-1, keepdims=True), 1e-5)
        current_data = np.round(input_data * scale).clip(-128, 127)

        for layer_info in self.quantized_model[:-1]:  # For all layers except the last one
            # print(f'layer: {layer_info["layer_type"], layer_info["layer_order"] }')

            if layer_info['layer_type'] == 'BitLinear':

                if len(current_data.shape) == 4:
                    # reshape from (batch_size, channels, height, width) to (batch_size, features)
                    current_data = current_data.reshape(current_data.shape[0], current_data.shape[1] * current_data.shape[2] * current_data.shape[3])

                weights = np.array(layer_info['quantized_weights'])
                conv = np.dot(current_data, weights.T)  # Matrix multiplication

                if layer_info['WScale']=='PerOutput':
                    scale = np.array(layer_info['quantized_scale']).transpose()
                    conv = conv * scale

                max = np.maximum(conv.max(axis=-1, keepdims=True), 1e-5)
                rescale = np.exp2(np.floor(np.log2(127.0 / max))) # Emulate normalization by shift as in C inference engine
                # rescale = 127.0 / np.maximum(conv.max(axis=-1, keepdims=True), 1e-5)   # Normalize to max 1.7 range
                current_data = np.round(conv * rescale).clip(0, 127)  # Quantize the output and ReLU

            elif layer_info['layer_type'] == 'BitConv2d':

                if len(current_data.shape) == 2:
                    # Reshape from (batch_size, features) to (batch_size, channels, height, width)
                    height = width = int(np.sqrt(current_data.shape[1] // layer_info['in_channels']))
                    current_data = current_data.reshape(current_data.shape[0], layer_info['in_channels'], height, width)

                kernel_size = layer_info['kernel_size'][0]  # Assuming square kernel
                groups = layer_info['groups']
                in_channels = layer_info['in_channels']
                out_channels = layer_info['out_channels']

                weights = np.array(layer_info['quantized_weights']).reshape(
                    out_channels, in_channels // groups, kernel_size, kernel_size)

                # print(f'weights: {weights.shape} data: {current_data.shape}')
                output = np.zeros((current_data.shape[0], layer_info['out_channels'],
                                current_data.shape[2] - kernel_size + 1, current_data.shape[3] - kernel_size + 1))

                # update the incoming and outgoing dimensions
                layer_info['incoming_x'] = current_data.shape[2]
                layer_info['incoming_y'] = current_data.shape[3]

                layer_info['outgoing_x'] = output.shape[2]
                layer_info['outgoing_y'] = output.shape[3]

                for g in range(groups):
                    for i in range(output.shape[2]):
                        for j in range(output.shape[3]):
                            patch = current_data[:, g*(in_channels//groups):(g+1)*(in_channels//groups),
                                                i:i+kernel_size, j:j+kernel_size]
                            group_weights = weights[g*(out_channels//groups):(g+1)*(out_channels//groups)]
                            output[:, g*(out_channels//groups):(g+1)*(out_channels//groups), i, j] = \
                                np.sum(patch[:, np.newaxis, :, :, :] * group_weights, axis=(2, 3, 4))

                # Apply ReLU and quantize
                output = np.maximum(output, 0)
                max_val = np.max(output, axis=(1, 2, 3), keepdims=True)
                # max_val = 256*127.0
                # print(max_val / 256)
                current_data = np.round(output * (127.0 / max_val)).clip(0, 127).astype(np.int8)

            elif layer_info['layer_type'] == 'MaxPool2d':
                pool_size = layer_info['kernel_size']
                stride = layer_info['stride']
                batch_size, channels, height, width = current_data.shape
                pooled_height = (height - pool_size) // stride + 1
                pooled_width = (width - pool_size) // stride + 1
                pooled_output = np.zeros((batch_size, channels, pooled_height, pooled_width), dtype=current_data.dtype)

                layer_info['incoming_x'] = current_data.shape[2]
                layer_info['incoming_y'] = current_data.shape[3]

                layer_info['outgoing_x'] = current_data.shape[2] // 2
                layer_info['outgoing_y'] = current_data.shape[2] // 2

                for i in range(pooled_height):
                    for j in range(pooled_width):
                        h_start = i * stride
                        h_end = h_start + pool_size
                        w_start = j * stride
                        w_end = w_start + pool_size
                        pooled_output[:, :, i, j] = np.max(current_data[:, :, h_start:h_end, w_start:w_end], axis=(2, 3))

                current_data = pooled_output          
                

        # no renormalization for the last layer
        weights = np.array(self.quantized_model[-1]['quantized_weights'])
        logits = np.dot(current_data, weights.T)  # Matrix multiplication

        # print(logits)  
        if self.quantized_model[-1]['WScale']=='PerOutput':
            scale = np.array(self.quantized_model[-1]['quantized_scale']).transpose()
            logits = logits * scale

        return logits