Skip to main content
`Ctrl`+`K`

# Getting started

Brevitas serves various types of users and end goals. With respect to defining quantized models, Brevitas supports two types of user flows:

- By hand, writing a quantized model using `brevitas.nn` quantized layers, possibly by modifying an original PyTorch floating-point model definition.

- Programmatically, by taking a floating-point model as input and automatically deriving a quantized model definition from it according to some user-defined criteria.

Once a quantize model is defined in either way, it can then be used as a starting point for either:

- PTQ (Post-Training Quantization), starting from a pretrained floating-point model.

- QAT (Quantization Aware Training), either training from scratch or finetuning a pretrained floating-point model.

- PTQ followed by QAT finetuning, to combine the best of both approaches.

## PTQ over hand or programmatically defined quantized models

Checkout how it’s done for ImageNet classification over torchvision or other hand-defined models with our example scripts here.

## Defining a quantized model with brevitas.nn layers

We consider quantization of a classic neural network, LeNet-5.

### Weights-only quantization, float activations and biases

Let’s say we are interested in assessing how well the model does at 4 bit weights for CIFAR10 classification. For the purpose of this tutorial we will skip any detail around how to perform training, as training a neural network with Brevitas is no different than training any other neural network in PyTorch.

`brevitas.nn` provides quantized layers that can be used in place of and/or mixed with traditional `torch.nn` layers. In this case then we import `brevitas.nn.QuantConv2d` and `brevitas.nn.QuantLinear` in place of their PyTorch variants, and we specify `weight_bit_width=4`. For relu and max-pool, we leverage the usual `torch.nn.ReLU` and `torch.nn.functional.max_pool2d`.

The result is the following:

    from torch import nn
    from torch.nn import Module
    import torch.nn.functional as F

    import brevitas.nn as qnn

    class QuantWeightLeNet(Module):
        def __init__(self):
            super(QuantWeightLeNet, self).__init__()
            self.conv1 = qnn.QuantConv2d(3, 6, 5, bias=True, weight_bit_width=4)
            self.relu1 = nn.ReLU()
            self.conv2 = qnn.QuantConv2d(6, 16, 5, bias=True, weight_bit_width=4)
            self.relu2 = nn.ReLU()
            self.fc1   = qnn.QuantLinear(16*5*5, 120, bias=True, weight_bit_width=4)
            self.relu3 = nn.ReLU()
            self.fc2   = qnn.QuantLinear(120, 84, bias=True, weight_bit_width=4)
            self.relu4 = nn.ReLU()
            self.fc3   = qnn.QuantLinear(84, 10, bias=True, weight_bit_width=4)

        def forward(self, x):
            out = self.relu1(self.conv1(x))
            out = F.max_pool2d(out, 2)
            out = self.relu2(self.conv2(out))
            out = F.max_pool2d(out, 2)
            out = out.reshape(out.reshape[0], -1)
            out = self.relu3(self.fc1(out))
            out = self.relu4(self.fc2(out))
            out = self.fc3(out)
            return out

    quant_weight_lenet = QuantWeightLeNet()

    # ... training ...

A neural network with 4 bits weights and floating-point activations can provide an advantage in terms of model storage, but it doesn’t provide any advantage in terms of compute, as the weights need to be converted to float at runtime first. In order to make more practical, we want to quantize activations too.

### Weights and activations quantization, float biases

We now quantize both weights and activations to 4 bits, while keeping biases in floating-point. In order to do so, we replace `torch.nn.ReLU` with `brevitas.nn.QuantReLU`, specifying `bit_width=4`. Additionally, in order to quantize the very first input, we introduce a `brevitas.nn.QuantIdentity` at the beginning of the network. The end result is the following:

    from torch.nn import Module
    import torch.nn.functional as F

    import brevitas.nn as qnn
    from brevitas.quant import Int8Bias as BiasQuant

    class QuantWeightActLeNet(Module):
        def __init__(self):
            super(QuantWeightActLeNet, self).__init__()
            self.quant_inp = qnn.QuantIdentity(bit_width=4)
            self.conv1 = qnn.QuantConv2d(3, 6, 5, bias=True, weight_bit_width=4)
            self.relu1 = qnn.QuantReLU(bit_width=4)
            self.conv2 = qnn.QuantConv2d(6, 16, 5, bias=True, weight_bit_width=4)
            self.relu2 = qnn.QuantReLU(bit_width=3)
            self.fc1   = qnn.QuantLinear(16*5*5, 120, bias=True, weight_bit_width=4)
            self.relu3 = qnn.QuantReLU(bit_width=4)
            self.fc2   = qnn.QuantLinear(120, 84, bias=True, weight_bit_width=4)
            self.relu4 = qnn.QuantReLU(bit_width=4)
            self.fc3   = qnn.QuantLinear(84, 10, bias=True)

        def forward(self, x):
            out = self.quant_inp(x)
            out = self.relu1(self.conv1(out))
            out = F.max_pool2d(out, 2)
            out = self.relu2(self.conv2(out))
            out = F.max_pool2d(out, 2)
            out = out.reshape(out.shape[0], -1)
            out = self.relu3(self.fc1(out))
            out = self.relu4(self.fc2(out))
            out = self.fc3(out)
            return out

    quant_weight_act_lenet = QuantWeightActLeNet()

    # ... training ...

Note a couple of things:

- By default `QuantReLU` is stateful, so there is a difference between instantiating one `QuantReLU` that is called multiple times, and instantiating multiple `QuantReLU`` that are each called once.

- `QuantReLU` first computes a relu function, and then quantizes its output. To take advantage of the fact that the output of relu is `>= 0` then, by default `QuantReLU` performs unsigned quantization, meaning in this case its output is `int4` data in `[0, 15]`.

- Quantized data in Brevitas is always represented in dequantized format, meaning that is represented within a float tensor. The output of QuantReLU then looks like a standard float torch Tensor, but it’s restricted to 16 different values (with 4 bits quantization). In order to get a more informative representation of quantized data, we need to set `return_quant_tensor=True`.

### Weights, activations, biases quantization

    from torch.nn import Module
    import torch.nn.functional as F

    import brevitas.nn as qnn
    from brevitas.quant import Int32Bias

    class QuantWeightActBiasLeNet(Module):
        def __init__(self):
            super(LowPrecisionLeNet, self).__init__()
            self.quant_inp = qnn.QuantIdentity(bit_width=4, return_quant_tensor=True)
            self.conv1 = qnn.QuantConv2d(3, 6, 5, bias=True, weight_bit_width=4, bias_quant=Int32Bias)
            self.relu1 = qnn.QuantReLU(bit_width=4, return_quant_tensor=True)
            self.conv2 = qnn.QuantConv2d(6, 16, 5, bias=True, weight_bit_width=4, bias_quant=Int32Bias)
            self.relu2 = qnn.QuantReLU(bit_width=4, return_quant_tensor=True)
            self.fc1   = qnn.QuantLinear(16*5*5, 120, bias=True, weight_bit_width=4, bias_quant=Int32Bias)
            self.relu3 = qnn.QuantReLU(bit_width=4, return_quant_tensor=True)
            self.fc2   = qnn.QuantLinear(120, 84, bias=True, weight_bit_width=4, bias_quant=Int32Bias)
            self.relu4 = qnn.QuantReLU(bit_width=4, return_quant_tensor=True)
            self.fc3   = qnn.QuantLinear(84, 10, bias=True, weight_bit_width=4, bias_quant=Int32Bias)

        def forward(self, x):
            out = self.quant_inp(x)
            out = self.relu1(self.conv1(out))
            out = F.max_pool2d(out, 2)
            out = self.relu2(self.conv2(out))
            out = F.max_pool2d(out, 2)
            out = out.reshape(out.shape[0], -1)
            out = self.relu3(self.fc1(out))
            out = self.relu4(self.fc2(out))
            out = self.fc3(out)
            return out

    quant_weight_act_bias_lenet = QuantWeightActBiasLeNet()

    # ... training ...

Compared to the previous scenario:

- We now set `return_quant_tensor=True` in every quantized activations to propagate a `QuantTensor` to the next layer. This informs each `QuantLinear` or `QuantConv2d` of how the input passed in has been quantized.

- A `QuantTensor` is just a tensor-like data structure providing metadata about how a tensor has been quantized, similar to a torch.qint dtype, but training friendly. Setting `return_quant_tensor=True` does not affect the way quantization is performed, it only changes the way the output is represented.

- We enable bias quantization by setting the Int32Bias quantizer. What it does is to perform bias quantization with ``bias_scale = input_scale * weight_scale`, as it commonly done across inference toolchains. This is why we have to set `return_quant_tensor=True`: each layer with `Int32Bias` can read the input scale from the `QuantTensor` passed in and use for bias quantization.

- `torch` operations that are algorithmically invariant to quantization, such as F.max_pool2d, can propagate QuantTensor through them without extra changes.

### Export to ONNX

Brevitas does not perform any low-precision acceleration on its own. For that to happen, the model need to be exported first to an inference toolchain through some intermediate representation like ONNX. One popular way to represent 8-bit quantization within ONNX is through the QDQ format. Brevitas extends QDQ to QCDQ, inserting a Clip node to represent quantization to <= 8 bits. We can then export all previous defined model to QCDQ. The interface of the export function matches the torch.onnx.export function, and accepts all its kwargs:

    from brevitas.export import export_onnx_qcdq
    import torch

    # Weight-only model
    export_onnx_qcdq(quant_weight_lenet, torch.randn(1, 3, 32, 32), export_path='4b_weight_lenet.onnx')

    # Weight-activation model
    export_onnx_qcdq(quant_weight_act_lenet, torch.randn(1, 3, 32, 32), export_path='4b_weight_act_lenet.onnx')

    # Weight-activation-bias model
    export_onnx_qcdq(quant_weight_act_bias_lenet, torch.randn(1, 3, 32, 32), export_path='4b_weight_act_bias_lenet.onnx')

## Where to go from here

Check out the Tutorials section for more information on things like ONNX export, quantized recurrent layers, quantizers, or a more detailed overview of the library in the TVMCon tutorial.