Convert a QONNX model to a Brevitas model

[Vylfred]

Hi everyone,

I'm working with Brevitas to run simulations and need to load a trained QONNX model into a new notebook using Brevitas/PyTorch. However, I've encountered an issue.

I've successfully loaded the weights for layers such as conv1, conv2, fc1, fc2, and fc3 from a LeNet-5 model. However, I'm having trouble with the scale factors. For instance, the scale of model.relu1.act_quant is 0.003921568859368563, whereas the corresponding scale in the QONNX model is 0.023734445. Here's what I tried: model.relu1.act_quant.scale().item() = weights['/relu1/act_quant/export_handler/Constant_output_0']. This results in the following error: SyntaxError: cannot assign to function call here. Maybe you meant '==' instead of '='?

I'm unsure how to update the scale parameters. How can I properly set the scale factor for the ReLU activation quantizer, for example?

Is there a specific method or approach for converting from ONNX to Brevitas that I might not be aware of?

Thanks in advance for any assistance!

Best,
Vylfrëd

[Giuseppe5]

Hello,
If the model is exported from Brevitas to QONNX, there is no need to manually set any scales, they are automatically picked up from the quantized model.

Could you post a script to reproduce your issue?

[Vylfred]

Hi,

Thank you for your previous response

I'm currently working with two Notebooks: one for training my model and saving it in QONNX format, and a second one where I need to reload the same model to perform fault simulations.

Here's the process I've followed for the second notebook (you will find it at the end of this response):

1. Loading the ONNX Model:

onnx_model = onnx.load(onnx_model_path)

weights = {}
for initializer in onnx_model.graph.initializer:
    weights[initializer.name] = onnx.numpy_helper.to_array(initializer)

2. Recreating the Model Architecture:

I recreated the same model architecture as in the first Notebook:

class QuantLeNet(nn.Module):
    def __init__(self):
        super(QuantLeNet, self).__init__()
        self.quant_input = qnn.QuantIdentity(bit_width=datasize, return_quant_tensor=True)
        self.conv1 = qnn.QuantConv2d(1, 6, 5, stride=1, padding=2, bias=False, weight_bit_width=datasize, bias_quant=Int8Bias)
        self.relu1 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = qnn.QuantConv2d(6, 16, 5, stride=1, bias=False, weight_bit_width=datasize, bias_quant=Int8Bias)
        self.relu2 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        self.fc1   = qnn.QuantLinear(16*5*5, 120, bias=False, weight_bit_width=datasize)
        self.relu3 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True)
        self.fc2   = qnn.QuantLinear(120, 84, bias=False, weight_bit_width=datasize)
        self.relu4 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True)
        self.fc3   = qnn.QuantLinear(84, 10, bias=False, weight_bit_width=datasize)

    def forward(self, x):
        x = self.quant_input(x)
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.fc2(x)
        x = self.relu4(x)
        x = self.fc3(x)
        return x

model = QuantLeNet()

3. Replacing Weights:

I successfully replaced all the weights from the QONNX model as follows:

with torch.no_grad(): 
    model.conv1.weight.copy_(torch.from_numpy(weights['/conv1/weight_quant/export_handler/Constant_1_output_0'] * weights['/conv1/weight_quant/export_handler/Constant_output_0']))
    model.conv2.weight.copy_(torch.from_numpy(weights['/conv2/weight_quant/export_handler/Constant_1_output_0'] * weights['/conv2/weight_quant/export_handler/Constant_output_0']))
    model.fc1.weight.copy_(torch.from_numpy(weights['/fc1/weight_quant/export_handler/Constant_1_output_0'] * weights['/fc1/weight_quant/export_handler/Constant_output_0']))
    model.fc2.weight.copy_(torch.from_numpy(weights['/fc2/weight_quant/export_handler/Constant_1_output_0'] * weights['/fc2/weight_quant/export_handler/Constant_output_0']))
    model.fc3.weight.copy_(torch.from_numpy(weights['/fc3/weight_quant/export_handler/Constant_1_output_0'] * weights['/fc3/weight_quant/export_handler/Constant_output_0']))

Issue: Replacing Scale Factors

The challenge I'm facing now is related to the scale factors. When comparing the scale factors between the new model and the QONNX model, they are different:

print(model.relu1.act_quant.scale().item()) # New model: 5.8493431424722075e-05
print(weights['/relu1/act_quant/export_handler/Constant_output_0']) # QONNX model: 0.023734445

I attempted to replace the scale factor from the QONNX model in the new model using the following methods:

• Direct Assignment (which failed with a syntax error):

model.relu1.act_quant.scale().item() = weights['/relu1/act_quant/export_handler/Constant_output_0']

Error: SyntaxError: cannot assign to function call here. Maybe you meant '==' instead of '='?

• Passing the Scale Factor in the Class Constructor:

self.relu1 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True, input_quant=weights['/relu1/act_quant/export_handler/Constant_output_0'])

Error: RuntimeError: The quantizer passed does not adhere to the quantization protocol.

• Using a Variable for the Scale Factor:

test_quant = model.relu1.act_quant.scale()

This approach also resulted in the same error.

Here is the full python code for the second notebook:

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import brevitas.nn as qnn

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

from brevitas.quant import Int8WeightPerTensorFixedPoint
from brevitas.quant import Int8ActPerTensorFloat
from brevitas.quant.scaled_int import Int8Bias
from brevitas.quant import Int8Bias
from brevitas.nn import QuantLinear
# from brevitas.export import export_finn_onnx
from brevitas.export import export_qonnx

import onnx
import onnxruntime
import onnxoptimizer
import numpy as np

torch.set_num_threads(1)


# # Load ONNX model

# In[2]:


# Load the ONNX model 
onnx_model_path = "8b_LeNet5_MNIST_qonnx.onnx"
onnx_model = onnx.load(onnx_model_path)

# Extract the layer weights from the ONNX model
weights = {}
for initializer in onnx_model.graph.initializer:
    weights[initializer.name] = onnx.numpy_helper.to_array(initializer)

#print(weights.keys())


# # Brevitas model

# In[3]:


num_bits = 8
datasize = num_bits


# In[4]:


class QuantLeNet(nn.Module):
    def __init__(self):
        super(QuantLeNet, self).__init__()
        self.quant_input = qnn.QuantIdentity(bit_width=datasize, return_quant_tensor=True)
        self.conv1 = qnn.QuantConv2d(1, 6, 5, stride=1, padding=2, bias=False, weight_bit_width=datasize, bias_quant=Int8Bias)
        self.relu1 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = qnn.QuantConv2d(6, 16, 5, stride=1, bias=False, weight_bit_width=datasize, bias_quant=Int8Bias)
        self.relu2 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        self.fc1   = qnn.QuantLinear(16*5*5, 120, bias=False, weight_bit_width=datasize)
        self.relu3 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True)
        self.fc2   = qnn.QuantLinear(120, 84, bias=False, weight_bit_width=datasize)
        self.relu4 = qnn.QuantReLU(bit_width=datasize, return_quant_tensor=True)
        self.fc3   = qnn.QuantLinear(84, 10, bias=False, weight_bit_width=datasize)

    def forward(self, x):
        x = self.quant_input(x) # 1
        x = self.conv1(x) # 2
        x = self.relu1(x) # 3
        x = self.pool1(x) # 4
        x = self.conv2(x) # 5
        x = self.relu2(x) # 6
        x = self.pool2(x) # 7
        x = self.flatten(x) # 8
        x = self.fc1(x) # 9
        x = self.relu3(x) # 10
        x = self.fc2(x) # 11
        x = self.relu4(x) # 12
        x = self.fc3(x) # 13
        return x
    
model = QuantLeNet()


# # Copy weights

# In[5]:


with torch.no_grad():
    model.conv1.weight.copy_(torch.from_numpy(weights['/conv1/weight_quant/export_handler/Constant_1_output_0'] * weights['/conv1/weight_quant/export_handler/Constant_output_0']))
    model.conv2.weight.copy_(torch.from_numpy(weights['/conv2/weight_quant/export_handler/Constant_1_output_0'] * weights['/conv2/weight_quant/export_handler/Constant_output_0']))
    model.fc1.weight.copy_(torch.from_numpy(weights['/fc1/weight_quant/export_handler/Constant_1_output_0'] * weights['/fc1/weight_quant/export_handler/Constant_output_0']))
    model.fc2.weight.copy_(torch.from_numpy(weights['/fc2/weight_quant/export_handler/Constant_1_output_0'] * weights['/fc2/weight_quant/export_handler/Constant_output_0']))
    model.fc3.weight.copy_(torch.from_numpy(weights['/fc3/weight_quant/export_handler/Constant_1_output_0'] * weights['/fc3/weight_quant/export_handler/Constant_output_0']))

dummy_input = torch.randn(1, 1, 28, 28)
output = model(dummy_input)


# # Evaluation

# In[6]:


import torch
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader

# Définir les transformations pour les données MNIST
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# Charger les datasets de test
test_dataset = MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# Fonction d'évaluation pour calculer l'accuracy
def evaluate_model(model, data_loader):
    model.eval()  # Mettre le modèle en mode évaluation
    correct = 0
    total = 0

    with torch.no_grad():
        for data, target in data_loader:
            output = model(data)
            _, predicted = torch.max(output, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

    accuracy = correct / total
    return accuracy

# Évaluer le modèle sur le dataset de test MNIST
accuracy = evaluate_model(model, test_loader)
print(f'Accuracy of the model on the MNIST test dataset: {accuracy * 100:.2f}%')

At this link, you can find the LeNet-5 onnx file : https://drive.google.com/file/d/1UmGubI5Xsq7orJWRLmZ7fJJ0dbwSzcD9/view?usp=sharing

• Additional Context:

Without modifying the scale factor, I observed a significant drop in accuracy (e.g., from ~80% to 40% on an AlexNet model with CIFAR-10). For a LeNet5 model, the accuracy drop is minimal, from 99% to 98%.

How can I correctly replace the scale factor in the new model with the value from the QONNX model?

Thanks for your time and assistance!
Vylfrëd

[nickfraser]

In short, QONNX is an export target from Brevitas and not designed to be a 2-way procedure. Similarly, PyTorch cannot import ONNX models (at least, without some external library). Is there any reason why you're not simply saving/loading the state_dict? Is there a reason why you're exporting to QONNX then reimporting into PyTorch/Brevitas?

If you insist on loading the scale factors from a trained QONNX format, then I suggest using a quantizer with a constant scale factor, and using those weights to initialise that constant value.