Converting from integer weights to original float weights

[khoaguin]

I have this model

from dependencies import value

import brevitas
import brevitas.nn as qnn
from brevitas.nn import QuantReLU, QuantIdentity, QuantLinear
from brevitas.core.bit_width import BitWidthImplType
from brevitas.core.quant import QuantType
from brevitas.core.restrict_val import FloatToIntImplType, RestrictValueType
from brevitas.core.scaling import ScalingImplType
from brevitas.core.zero_point import ZeroZeroPoint
from brevitas.inject import ExtendedInjector
from brevitas.quant.solver import ActQuantSolver, WeightQuantSolver

import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.nn.utils import prune
from torch.utils.data import DataLoader, TensorDataset
import torch.nn.functional as F


class CommonQuant(ExtendedInjector):
    bit_width_impl_type = BitWidthImplType.CONST
    scaling_impl_type = ScalingImplType.CONST
    restrict_scaling_type = RestrictValueType.FP
    zero_point_impl = ZeroZeroPoint
    float_to_int_impl_type = FloatToIntImplType.ROUND
    scaling_per_output_channel = False
    narrow_range = True
    signed = True

    @value
    def quant_type(bit_width):
        if bit_width is None:
            return QuantType.FP
        if bit_width == 1:
            return QuantType.BINARY
        return QuantType.INT


class CommonWeightQuant(CommonQuant, WeightQuantSolver):
    scaling_const = 1.0
    signed = False


def _make_qlinear_fx(in_features, out_features, bias=False, weight_bit_width=2):
    return qnn.QuantLinear(
        in_features=in_features,
        out_features=out_features,
        bias=bias,
        bias_quant=None,
        narrow_range=True,
        weight_bit_width=weight_bit_width,
        weight_quant=CommonWeightQuant,
        # return_quant_tensor=True
    )

IN_FEAT = 28*28
N_HIDDEN = 128
OUT_FEAT = 10

class QATPrunedSimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.pruned_layers = set()
        self.quant_inp = qnn.QuantIdentity(
            bit_width=3, act_quant=None, 
            # return_quant_tensor=True
        )
        self.fc1 = _make_qlinear_fx(IN_FEAT, N_HIDDEN)
        self.fc2 = _make_qlinear_fx(N_HIDDEN, N_HIDDEN)
        self.fc3 = _make_qlinear_fx(N_HIDDEN, OUT_FEAT)
        for m in self.modules():
            if isinstance(m, qnn.QuantLinear):
                torch.nn.init.uniform_(m.weight.data, 0, 1)

    def forward(self, x):
        x = self.quant_inp(x)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = QATPrunedSimpleNet()

After training the network, if I do

int_weight = model.fc1.quant_weight().int()
print(f"torch.unique(int_weight) = {torch.unique(int_weight)}")
zero_point = model.fc1.weight_quant.zero_point()
print(f"zero_point = {zero_point}")
scale = model.fc1.weight_quant.scale()
print(f"scale = {scale}")
quant_weight_manually = (int_weight - zero_point) * scale
print(f"{torch.unique(quant_weight_manually) = }")
is_close = model.fc1.quant_weight().value.isclose(quant_weight_manually).all().item()
print(f"quant_weight() is close to quant_weight_manually: {is_close}")

I will get

torch.unique(int_weight) = tensor([0, 1, 2], dtype=torch.uint8)
zero_point = 0.0
scale = 0.5
torch.unique(quant_weight_manually) = tensor([0.0000, 0.5000, 1.0000])
quant_weight() is close to quant_weight_manually: True

However, I want a way to convert between the original floating point weight, e.g. model.fc1.weight and the integer weight (model.fc1.quant_weight().int()). How would I do so?

[Giuseppe5]

Hi, thanks for your question.

I must say, I'm not I follow what you mean with "convert from the original weights and the integer weight". Could you elaborate a bit more?

[khoaguin]

Original weights as in model.fc1.weight (PyTorch's weights)

[Giuseppe5]

What kind of conversion are you trying to achieve?

[khoaguin]

From model.fc1.weight to model.fc1.quant_weight().int()

[Giuseppe5]

When you call model.fc1.quant_weight().int(), what happens under the hood is that the original weights of the model (model.fc1.weight) are dynamically quantized and converted then to integer.
We don't store the quantized weights usually, unless in special cases with context managers and specific flags

[khoaguin]

So my question is a function to convert from model.fc1.quant_weight().int() to model.fc1.weight and back. Can we do that?

[Giuseppe5]

Convert from model.fc1.quant_weight().int() to model.fc1.weight is not possible since the quantized operation introduces rounding and clipping that can't be reverted easily.

Might I ask what you're trying to achieve with this function? That could help me answering your question better

[khoaguin]

Essentially I want to do FL with Brevitas models, but the averaging step needs to be done in integers (with rounding of course). Then I want to evaluate using pure integer values. Then, I will convert the averaged model to float weights so I can load it back to PyTorch model and train again in the next round. Is this possible?

[Giuseppe5]

It depends on what averaging you are doing. If the scale factors of the integer tensors you are averaging are all the same, then you can: apply the dequantization formula to go back to floating point and load that back to pytorch

deq_w = (quant_w - zero_point) * scale

Scale and zero_point are quantization metadata that you can find within the QuantTensor.

If the scale factors of the tensors you are averaging are different, then I'm not entirely sure how to proceed.