ml-sdk/use_case_examples/cifar/cifar_brevitas_finetuning/FromImageNetToCifar.ipynb

Part 1: From ImageNet to CIFAR with Transfer Learning

This notebook may be skipped if the user already has a pre-trained floating point CIFAR-10 / CIFAR-100 model. In what follows, we show you how we used Transfer Learning to train a convolutional floating point neural network.

Nowadays, whether or not you have time and computational resources, very few people train a full convolutional network on a huge data-set from scratch, e.g., ImageNet, which has 1.2 million images and 1000 categories. In practice, we use Transfer learning a very efficient approach in Machine Learning that aims to save both time and computing power. Transfer learning refers to a set of methods that allow the reuse of knowledge gained from solving a given task and applying it to a different but related task.

Transfer Learning provides a variety of techniques, such as:

• fine-tuning: instead of a random initialization of the weights, it's possible to reuse the pre-trained weights as a starting point
• fixed feature extractor: the pre-trained model is used as a feature extractor and only the output model is fine-tuned to achieve the targeted task. In other words, the final fully connected layer is replaced by a new one with the appropriate output shape. We freeze all the layers except this new one, that is trained.

In our use case, we'll combine the two methods mentioned above.

Import the required packages.

In [1]:

import random
import warnings
from copy import deepcopy

import torch
from cifar_utils import get_dataloader, plot_dataset, plot_history, torch_inference, train
from models import Fp32VGG11

warnings.filterwarnings("ignore")

Settings

In [2]:

dataset_name = "CIFAR_100"

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

param_c10 = {
    "output_size": 10,
    "batch_size": 128,
    "training": "fp32",
    "dataset_name": "CIFAR_10",
    "criterion": torch.nn.CrossEntropyLoss(),
    "accuracy_test": [],
    "accuracy_train": [],
    "loss_test_history": [],
    "loss_train_history": [],
    "dir": "./checkpoints/CIFAR_10",
    "seed": 727,
}
param_c100 = {
    "output_size": 100,
    "batch_size": 128,
    "training": "fp32",
    "dataset_name": "CIFAR_100",
    "criterion": torch.nn.CrossEntropyLoss(),
    "accuracy_test": [],
    "accuracy_train": [],
    "loss_test_history": [],
    "loss_train_history": [],
    "dir": "./checkpoints/CIFAR_100",
    "seed": 727,
}

# In this tutorial, we present the results of CIFAR-100.
if dataset_name == "CIFAR_100":
    param = param_c100
else:
    # If you want to use it for CIFAR-10, set `dataset_name` to "CIFAR_100"
    param = param_c10

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

Device Type: cuda:0

Data-set

The CIFAR data-sets contain 32*32 colored images.

In [3]:

# Load CIFAR-100 or CIFAR-10 data-set according to `dataset_name`.
train_loader, test_loader = get_dataloader(param=param)

# Let’s visualize `n` images from CIFAR data-set.
plot_dataset(test_loader, param)

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).

Model's architecture

As this network is meant to run in FHE, we apply several changes on the original VGG11 architecture in order to speed up the FHE execution and therefore make it more user-friendly :

• replacing the MaxPool2d by the AvgPool2d, because in the current version of TorusML MaxPool2d isn't available yet
• changing the kernel and stride size in some pooling layers, because we have chosen to keep the initial input size of 3*32*32 instead of the recommended input size of 3*224*224
• changing the kernel size from 7 to 1 in the AdaptiveAvgPool2d to further reduce the image's size
• removing the classification part of the original VGG11 architecture to speed up computation

The model architecture is detailed here

Load VGG-11 pre-trained weights

In [4]:

torch.manual_seed(param["seed"])
random.seed(param["seed"])

# Instantiation of the custom VGG-11 network.
fp32_vgg = Fp32VGG11(param["output_size"]).to(device)

# Loading the pre-trained VGG-11 weights from torch.hub.
pretrained_weights = torch.hub.load(
    "pytorch/vision:v0.10.0",
    "vgg11",
    pretrained=True,
).state_dict()

# Caution:
# The `fp32_vgg.state_dict()` respects the same schema for the convolutional, ReLU and pooling
# layers (order/shape/naming) as the `pretrained_weights` dict, that's why it doesn't throw an error
fp32_vgg.load_state_dict(deepcopy(pretrained_weights), strict=False)
# We got an IncompatibleKeys warning because
# we deleted the classification layers of the original VGG-11 network.

Out [4]:

Using cache found in /home/stoiana/.cache/torch/hub/pytorch_vision_v0.10.0

_IncompatibleKeys(missing_keys=['final_layer.weight', 'final_layer.bias'], unexpected_keys=['classifier.0.weight', 'classifier.0.bias', 'classifier.3.weight', 'classifier.3.bias', 'classifier.6.weight', 'classifier.6.bias'])

Training

For the training part, we combine 2 methods to fine-tune our model.

1. Fine-tuning the final layer: here, the weights are initialized using VGG-11 from torch-hub. We freeze the convolutional layers except the last fully connected final layer, which we train. This approach prevents abrupt gradient propagation.

2. Fine-tuning all the layers: in this case, we unfreeze the convolutional layers and fine-tune the entire network.

Fine-tuning the final layer

We freeze all the layers except the final fully connected layer, that we fine-tune.

In [5]:

# We freeze all the layers.
for p in list(fp32_vgg.parameters()):
    p.requires_grad = False

# Set the `requires_grad` of the last layer at `True` to fine-tune it.
fp32_vgg.final_layer.weight.requires_grad = True
fp32_vgg.final_layer.bias.requires_grad = True

In [6]:

if dataset_name == "CIFAR_100":
    param["lr"] = 0.06
    param["epochs"] = 10
    param["gamma"] = 0.01
    param["milestones"] = [1]

elif dataset_name == "CIFAR_10":
    param["lr"] = 0.1
    param["epochs"] = 5
    param["gamma"] = 0.1
    param["milestones"] = [1, 3]

fp32_vgg = train(fp32_vgg, train_loader, test_loader, param, device=device)

Epoch  0: Train loss = 2.5821 VS Test loss = 2.2989 - Accuracy train: 0.3639 VS Accuracy test: 0.4441
Epoch  1: Train loss = 2.0172 VS Test loss = 2.0736 - Accuracy train: 0.4741 VS Accuracy test: 0.4776
Epoch  2: Train loss = 1.9677 VS Test loss = 2.0398 - Accuracy train: 0.4868 VS Accuracy test: 0.4820
Epoch  3: Train loss = 1.9539 VS Test loss = 2.0230 - Accuracy train: 0.4887 VS Accuracy test: 0.4854
Epoch  4: Train loss = 1.9427 VS Test loss = 2.0072 - Accuracy train: 0.4941 VS Accuracy test: 0.4889
Epoch  5: Train loss = 1.9285 VS Test loss = 1.9941 - Accuracy train: 0.4978 VS Accuracy test: 0.4898
Epoch  6: Train loss = 1.9226 VS Test loss = 1.9830 - Accuracy train: 0.4960 VS Accuracy test: 0.4928
Epoch  7: Train loss = 1.9159 VS Test loss = 1.9764 - Accuracy train: 0.4965 VS Accuracy test: 0.4936
Epoch  8: Train loss = 1.9075 VS Test loss = 1.9681 - Accuracy train: 0.5017 VS Accuracy test: 0.4944
Epoch  9: Train loss = 1.9019 VS Test loss = 1.9612 - Accuracy train: 0.5023 VS Accuracy test: 0.4956
100%|██████████| 10/10 [10:37<00:00, 63.72s/it]

Fine-tune all layers

Now, unfreeze the convolution and fine-tune all the layers.

In [7]:

# Secondly, fine-tuning all the layers.
for p in list(fp32_vgg.parameters()):
    p.requires_grad = True

if dataset_name == "CIFAR_100":
    param["lr"] = 0.0006
    param["epochs"] = 7
    param["gamma"] = 0.1
    param["milestones"] = [6]

elif dataset_name == "CIFAR_10":
    param["lr"] = 0.0006
    param["epochs"] = 5
    param["gamma"] = 0.1
    param["milestones"] = [4]

fp32_vgg = train(fp32_vgg, train_loader, test_loader, param, device=device)

Epoch  0: Train loss = 2.3488 VS Test loss = 1.7114 - Accuracy train: 0.3919 VS Accuracy test: 0.5232
Epoch  1: Train loss = 1.5548 VS Test loss = 1.4423 - Accuracy train: 0.5648 VS Accuracy test: 0.5929
Epoch  2: Train loss = 1.2724 VS Test loss = 1.3197 - Accuracy train: 0.6374 VS Accuracy test: 0.6196
Epoch  3: Train loss = 1.0783 VS Test loss = 1.3551 - Accuracy train: 0.6842 VS Accuracy test: 0.6207
Epoch  4: Train loss = 0.9087 VS Test loss = 1.3401 - Accuracy train: 0.7301 VS Accuracy test: 0.6345
Epoch  5: Train loss = 0.7783 VS Test loss = 1.3401 - Accuracy train: 0.7665 VS Accuracy test: 0.6491
Epoch  6: Train loss = 0.3853 VS Test loss = 1.2485 - Accuracy train: 0.8849 VS Accuracy test: 0.6982
100%|██████████| 7/7 [09:43<00:00, 83.37s/it]

In [8]:

acc = torch_inference(fp32_vgg, data=test_loader, device=device)
print(f"With {dataset_name}: top-1 accuracy = {acc * 100:2.3f}%")

With CIFAR_100: top-1 accuracy = 69.822%

Let's visualize the accuracy and loss of the training and test sets.

In [9]:

plot_history(param)

The methodology is the same for CIFAR-10, re-run the notebook with dataset_name == CIFAR_10.

The model weights are also saved here and we get a model with the following performance :

In [10]:

plot_history(param_c10, load=True)

The sharp fluctuations are due to the change in the learning rate.

Conclusion

With such a simple architecture using AvgPool2d, 8 convolutional layers and 1 fully connected network, we achieve:

Data-set	Top-1 Accuracy	SOTA in 2022
CIFAR-10	$90.504$%	$99.612$%
CIFAR-100	$70.042$%	$96.808$%

In this notebook, we have shown how we obtained the weights needed to solve CIFAR-10 or CIFAR-100 classification task. These weights will be reused in the next tutorial. Where we'll use Brevitas to quantize these models as fully homomorphic encryption operates only over integer values.