torus-ml/use_case_examples/cifar/cifar_brevitas_finetuning/CifarQuantizationAwareTraining.ipynb

Part 2: Quantization Aware Training with CIFAR

In this notebook, we explain the constraints that fully homomorphic encryption (FHE) imposes on neural networks and how to quantize a custom pretrained neural network to make in compatible with FHE.

In order to make inference of neural networks work with FHE, two main constraints must be respected :

• Quantization: the input data and model's weights and activations must be integers. In Torus ML, this step is performed automatically for built-in models. In the case of Neural Networks, this is done using techniques such as Quantized Aware Training (QAT)through the use of Brevitas

• Precision limitation: The current version of Torus ML has a precision limitation of $16$-bits, i.e., it doesn't support encrypted values larger than $16$-bits integers

This notebook shows how to perform quantization aware training using Brevitas to solve a classification task on CIFAR-10 and CIFAR-100.

Import the required packages.

In [1]:

import torus.compiler
import torch
from cifar_utils import (
    fhe_compatibility,
    get_dataloader,
    mapping_keys,
    plot_baseline,
    plot_dataset,
    torch_inference,
    train,
)

# As we follow the same methodology for quantization aware training for CIFAR-10 and CIFAR-100.
# Let's import some generic functions.
from models import QuantVGG11
from torchvision import datasets

From torch to Brevitas network

Brevitas offers a quantized version of almost all torch layers. So, to quantize an existing torch model, you need to convert each torch layer to its Brevitas equivalent and provide some extra quantization parameters, like :

• bit_width: quantization bits for activations
• weight_bit_width: quantization bits for weights
• weight_quant : quantization protocol for the weights
• act_quant : quantization protocol for the activations

To quantize the input data, you have to add a very first QuantIdentity layer. It's also possible to use un-variate Pytorch layers mixed in with Brevitas quantized layers.

For instance with Brevitas, the Fp32VGG11 shown in the first notebook becomes QuantVGG11.

Settings

In [2]:

bit = 5

device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Device Type: {device}")

Device Type: cpu

Torus ML also supports a CUDA-enabled backend. To set it up, follow the instructions in the official guide for installing the GPU-enabled Torus compiler.

In [ ]:

compilation_device = "cuda" if torus.compiler.check_gpu_available() else "cpu"

print(f"Is GPU enabled: {torus.compiler.check_gpu_enabled()}")
print(f"Is GPU available: {torus.compiler.check_gpu_available()}")

1. Quantization Aware Training (QAT) with CIFAR-100

In [3]:

param_c100 = {
    "output_size": 100,
    "batch_size": 128,
    "training": "quant",
    "dataset_name": "CIFAR_100",
    "dataset": datasets.CIFAR100,
    "mean": [0.485, 0.456, 0.406],
    "std": [0.229, 0.224, 0.225],
    "lr": 6e-5,
    "seed": 727,
    "epochs": 10,
    "gamma": 0.01,
    "milestones": [3, 5],
    "criterion": torch.nn.CrossEntropyLoss(),
    "accuracy_test": [],
    "accuracy_train": [],
    "loss_test_history": [],
    "loss_train_history": [],
    "dir": "./checkpoints/CIFAR_100",
    "pre_trained_path": "fp32/CIFAR_100_fp32_state_dict.pt",
}

CIFAR-100 data-set

In [4]:

# Load CIFAR-100 data-set.
train_loader_c100, test_loader_c100 = get_dataloader(param=param_c100)

# Let’s visualize `n` images from CIFAR-100 data-set.
plot_dataset(test_loader_c100, param_c100)

Files already downloaded and verified
Files already downloaded and verified

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

Quantized model

In [5]:

quant_vgg = QuantVGG11(bit=bit, output_size=param_c100["output_size"])

In this case, there are two possible scenarios: either we initialize the model's weights with those of the pre-trained model or we train the model from scratch. In this notebook, we have opted to use the first strategy. Please refer to the first notebook for additional details on how we managed to get the fp32 weights.

In [6]:

checkpoint = torch.load(
    f"{param_c100['dir']}/{param_c100['pre_trained_path']}", map_location=device
)

Mapping keys

Caution: to load a pre-entrained fp32 state_dict into its equivalent quantized network it is necessary to have the same network architecture and same layer namings, otherwise it will raise an error.

The shift in the layer names in state_dict is due to the addition of QuantIdentity layers, that are necessary to ensure that the intermediate values of the model are all properly quantized.

In [7]:

quant_vgg = mapping_keys(checkpoint, quant_vgg, device)

Accuracy before QAT

In [9]:

# It's better than a random classifier.

acc_before_ft = torch_inference(quant_vgg, test_loader_c100, device=device)
param_c100["accuracy_test"].append(acc_before_ft)

print(f"Top 1 accuracy before fine-tuning = {acc_before_ft * 100:.4f}%")

Top 1 accuracy before fine-tuning = 7.3217%

FHE compatibility

It's a good practice to first check the FHE-compatibility and whether homomorphic inference is achievable or not before moving on.

To use this FHE simulation, the user needs to provide a representative set of data in order to determine the maximum bit-width within the circuit.

In [10]:

# The user can either provide the entire train data-set or a smaller but representative subset.
# As each batch is shuffled and contains 128 samples, it's a potential subset.

data_calibration, _ = next(iter(train_loader_c100))

qmodel = fhe_compatibility(quant_vgg, data_calibration, device=compilation_device)

print(
    f"With {param_c100['dataset_name']}, the maximum bit-width in the circuit = "
    f"{qmodel.fhe_circuit.graph.maximum_integer_bit_width()}"
)

Fine-tuning using Quantization aware training (QAT) with Brevitas

For the fine-tuning approach, it is recommended to start with a low learning rate.

In [10]:

quant_vgg = train(quant_vgg, train_loader_c100, test_loader_c100, param_c100, device=device)

Epoch  0: Train loss = 0.5070 VS Test loss = 1.3996 - Accuracy train: 0.8428 VS Accuracy test: 0.6791
Epoch  1: Train loss = 0.3827 VS Test loss = 1.4884 - Accuracy train: 0.8800 VS Accuracy test: 0.6805
Epoch  2: Train loss = 0.3163 VS Test loss = 1.5556 - Accuracy train: 0.8996 VS Accuracy test: 0.6873
Epoch  3: Train loss = 0.2609 VS Test loss = 1.5615 - Accuracy train: 0.9170 VS Accuracy test: 0.6855
Epoch  4: Train loss = 0.2505 VS Test loss = 1.5604 - Accuracy train: 0.9189 VS Accuracy test: 0.6877
Epoch  5: Train loss = 0.2497 VS Test loss = 1.5592 - Accuracy train: 0.9218 VS Accuracy test: 0.6889
Epoch  6: Train loss = 0.2498 VS Test loss = 1.5576 - Accuracy train: 0.9219 VS Accuracy test: 0.6882
Epoch  7: Train loss = 0.2442 VS Test loss = 1.5540 - Accuracy train: 0.9222 VS Accuracy test: 0.6894
Epoch  8: Train loss = 0.2434 VS Test loss = 1.5576 - Accuracy train: 0.9228 VS Accuracy test: 0.6879
Epoch  9: Train loss = 0.2504 VS Test loss = 1.5537 - Accuracy train: 0.9205 VS Accuracy test: 0.6897
100%|██████████| 10/10 [11:03<00:00, 66.38s/it]

In [11]:

plot_baseline(param_c100, test_loader_c100, device)

100%|██████████| 78/78 [00:02<00:00, 32.94it/s]

Running time: 0.0002 sec.

As we can see, convergence occurs very quickly, in just one epoch we nearly reached the accuracy of the pre-trained model in fp32. The fine-tuning strategy appears to be quite efficient to save time and computating power.

If you choose to train the model from scratch, you'll need more epochs and a thorough understanding of the Brevitas framework to achieve good results.

2. Quantization Aware Training (QAT) with CIFAR-10

Since the methodology is the same, we do exactly the same thing but this time with the CIFAR-10 hyper-parameters.

In [12]:

param_c10 = {
    "output_size": 10,
    "batch_size": 128,
    "training": "quant",
    "dataset_name": "CIFAR_10",
    "dataset": datasets.CIFAR10,
    "std": [0.247, 0.243, 0.261],
    "mean": [0.4914, 0.4822, 0.4465],
    "seed": 42,
    "lr": 1e-5,
    "epochs": 9,
    "gamma": 0.1,
    "milestones": [4],
    "criterion": torch.nn.CrossEntropyLoss(),
    "accuracy_test": [],
    "accuracy_train": [],
    "loss_test_history": [],
    "loss_train_history": [],
    "dir": "./checkpoints/CIFAR_10",
    "pre_trained_path": "fp32/CIFAR_10_fp32_state_dict.pt",
}

CIFAR-10 data-set

In [13]:

# Load CIFAR-10 data-set.

train_loader_c10, test_loader_c10 = get_dataloader(param=param_c10)

# Let’s visualize `n` images from CIFAR-100 data-set.

plot_dataset(test_loader_c10, param_c10)

Files already downloaded and verified
Files already downloaded and verified

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

Quantized model

In [14]:

# Instantiate our custom vgg11 network.
quant_vgg = QuantVGG11(bit=bit, output_size=param_c10["output_size"]).to(device)

checkpoint = torch.load(f"{param_c10['dir']}/{param_c10['pre_trained_path']}", map_location=device)
# checkpoint = torch.load("./checkpoints/CIFAR_10/fp32/CIFAR10_fp32_weights_acc_0.9013.pt")

# Mapping keys.
quant_vgg = mapping_keys(checkpoint, quant_vgg, device)

Accuracy before QAT

In [15]:

# It's better than a random classifier.

acc_before_ft = torch_inference(quant_vgg, test_loader_c10, device=device)
param_c10["accuracy_test"].append(acc_before_ft)

print(f"Top-1 accuracy before fine-tuning = {acc_before_ft * 100:.4f}%")

100%|██████████| 78/78 [00:05<00:00, 15.04it/s]

Running time: 0.0005 sec.
Top-1 accuracy before fine-tuning = 29.3570%

FHE-compatibility

In [16]:

# Check the FHE-compatibility.
data, _ = next(iter(train_loader_c10))

qmodel = fhe_compatibility(quant_vgg, data, device=compilation_device)

print(
    f"With {param_c10['dataset_name']}, the circuit has a maximum bit-width of "
    f"{qmodel.fhe_circuit.graph.maximum_integer_bit_width()}-bits."
)

With CIFAR_10, the circuit has a maximum bit-width of 14-bits.

Fine-tuning using Quantization aware training (QAT) with Brevitas

In [17]:

# Quantization Aware Training.
quant_vgg = train(quant_vgg, train_loader_c10, test_loader_c10, param_c10, device=device)

Epoch  0: Train loss = 0.2612 VS Test loss = 0.3473 - Accuracy train: 0.9084 VS Accuracy test: 0.8895
Epoch  1: Train loss = 0.2136 VS Test loss = 0.3421 - Accuracy train: 0.9252 VS Accuracy test: 0.8902
Epoch  2: Train loss = 0.1629 VS Test loss = 0.3267 - Accuracy train: 0.9443 VS Accuracy test: 0.8986
Epoch  3: Train loss = 0.1578 VS Test loss = 0.3236 - Accuracy train: 0.9467 VS Accuracy test: 0.8999
Epoch  4: Train loss = 0.1521 VS Test loss = 0.3227 - Accuracy train: 0.9470 VS Accuracy test: 0.9006
Epoch  5: Train loss = 0.1520 VS Test loss = 0.3185 - Accuracy train: 0.9472 VS Accuracy test: 0.9042
100%|██████████| 6/6 [07:19<00:00, 73.27s/it]

In [18]:

plot_baseline(param_c10, test_loader_c10, device)

After 3 epochs, we reached the accuracy of the pre-trained model in fp32.

Conclusion

By initializing the Brevitas FHE-friendly network with pre-trained fp32 weights and using quantization aware training approach we can see that the convergence is extremely fast, after one epoch we nearly catch up with the baseline in fp32.

However, the quantization process could cause a slight decrease in performance.

Data-set	Top1-accuracy with Brevitas	Top1-accuracy with Pytorch
CIFAR-10	$90.42$%	$90.11$%
CIFAR-100	$68.46$%	$70.43$%

In this notebook, we have explained why we need to quantize a model, how to check if a network is FHE-compliant and how to quantize an fp32 pre-trained model using Quantization Aware Training (QAT) with Bravitas.

In the third notebook, we will show how to evaluate a quantized model in FHE.