nncf_pytorch_8bit_quantization_with_openvino_export.py

1import torch
2import torchvision
3import nncf  # Neural Network Compression Framework
4import openvino as ov
5
6# 1. Prepare the model
7model = torchvision.models.resnet18(pretrained=True)
8model.eval()
9
10# 2. Prepare the dataset for calibration
11# In this example, we use a dummy dataset
12def transform(x):
13    return x
14
15dataset = torchvision.datasets.FakeData(size=10, image_size=(3, 224, 224), transform=torchvision.transforms.ToTensor())
16loader = torch.utils.data.DataLoader(dataset, batch_size=1)
17
18def transform_fn(data_item):
19    images, _ = data_item
20    return images
21
22# Wrap the loader into nncf.Dataset
23calibration_dataset = nncf.Dataset(loader, transform_fn)
24
25# 3. Run quantization
26quantized_model = nncf.quantize(model, calibration_dataset)
27
28# 4. (Optional) Export to OpenVINO Intermediate Representation (IR)
29dummy_input = torch.randn(1, 3, 224, 224)
30ov_model = ov.convert_model(quantized_model, example_input=dummy_input)
31ov.save_model(ov_model, "quantized_resnet18.xml")
32
33print("Quantization finished. Model saved to quantized_resnet18.xml")