GitHub - openvinotoolkit/nncf: Neural Network Compression Framework for enhanced OpenVINO™ inference (original) (raw)

Neural Network Compression Framework (NNCF) provides a suite of post-training and training-time algorithms for optimizing inference of neural networks in OpenVINO™ with a minimal accuracy drop.

NNCF is designed to work with models from PyTorch, TorchFX, TensorFlow, ONNX and OpenVINO™.

NNCF provides samples that demonstrate the usage of compression algorithms for different use cases and models. See compression results achievable with the NNCF-powered samples on the NNCF Model Zoo page.

The framework is organized as a Python* package that can be built and used in a standalone mode. The framework architecture is unified to make it easy to add different compression algorithms for both PyTorch and TensorFlow deep learning frameworks.

Key Features

Post-Training Compression Algorithms

Training-Time Compression Algorithms

• Automatic, configurable model graph transformation to obtain the compressed model.

  NOTE: Limited support for TensorFlow models. Only models created using Sequential or Keras Functional API are supported.

• Common interface for compression methods.
• GPU-accelerated layers for faster compressed model fine-tuning.
• Distributed training support.
• Git patch for prominent third-party repository (huggingface-transformers) demonstrating the process of integrating NNCF into custom training pipelines.
• Exporting PyTorch compressed models to ONNX* checkpoints and TensorFlow compressed models to SavedModel or Frozen Graph format, ready to use with OpenVINO™ toolkit.
• Support for Accuracy-Aware model training pipelines via the Adaptive Compression Level Training and Early Exit Training.

Documentation

This documentation covers detailed information about NNCF algorithms and functions needed for the contribution to NNCF.

The latest user documentation for NNCF is available here.

NNCF API documentation can be found here.

Usage

Post-Training Quantization

The NNCF PTQ is the simplest way to apply 8-bit quantization. To run the algorithm you only need your model and a small (~300 samples) calibration dataset.

OpenVINO is the preferred backend to run PTQ with, while PyTorch, TensorFlow, and ONNX are also supported.

OpenVINO

import nncf import openvino as ov import torch from torchvision import datasets, transforms

Instantiate your uncompressed model

model = ov.Core().read_model("/model_path")

Provide validation part of the dataset to collect statistics needed for the compression algorithm

val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()])) dataset_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1)

Step 1: Initialize transformation function

def transform_fn(data_item): images, _ = data_item return images

Step 2: Initialize NNCF Dataset

calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)

Step 3: Run the quantization pipeline

quantized_model = nncf.quantize(model, calibration_dataset)

PyTorch

import nncf import torch from torchvision import datasets, models

Instantiate your uncompressed model

model = models.mobilenet_v2()

Provide validation part of the dataset to collect statistics needed for the compression algorithm

val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()])) dataset_loader = torch.utils.data.DataLoader(val_dataset)

Step 1: Initialize the transformation function

def transform_fn(data_item): images, _ = data_item return images

Step 2: Initialize NNCF Dataset

calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)

Step 3: Run the quantization pipeline

quantized_model = nncf.quantize(model, calibration_dataset)

NOTE If the Post-Training Quantization algorithm does not meet quality requirements you can fine-tune the quantized pytorch model. You can find an example of the Quantization-Aware training pipeline for a pytorch model here.

TorchFX

import nncf import torch.fx from torchvision import datasets, models

Instantiate your uncompressed model

model = models.mobilenet_v2()

Provide validation part of the dataset to collect statistics needed for the compression algorithm

val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()])) dataset_loader = torch.utils.data.DataLoader(val_dataset)

Step 1: Initialize the transformation function

def transform_fn(data_item): images, _ = data_item return images

Step 2: Initialize NNCF Dataset

calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)

Step 3: Export model to TorchFX

input_shape = (1, 3, 224, 224) fx_model = torch.export.export_for_training(model, args=(ex_input,)).module()

or

fx_model = torch.export.export(model, args=(ex_input,)).module()

Step 4: Run the quantization pipeline

quantized_fx_model = nncf.quantize(fx_model, calibration_dataset)

TensorFlow

import nncf import tensorflow as tf import tensorflow_datasets as tfds

Instantiate your uncompressed model

model = tf.keras.applications.MobileNetV2()

Provide validation part of the dataset to collect statistics needed for the compression algorithm

val_dataset = tfds.load("/path", split="validation", shuffle_files=False, as_supervised=True)

Step 1: Initialize transformation function

def transform_fn(data_item): images, _ = data_item return images

Step 2: Initialize NNCF Dataset

calibration_dataset = nncf.Dataset(val_dataset, transform_fn)

Step 3: Run the quantization pipeline

quantized_model = nncf.quantize(model, calibration_dataset)

ONNX

import onnx import nncf import torch from torchvision import datasets

Instantiate your uncompressed model

onnx_model = onnx.load_model("/model_path")

Provide validation part of the dataset to collect statistics needed for the compression algorithm

val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()])) dataset_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1)

Step 1: Initialize transformation function

input_name = onnx_model.graph.input[0].name def transform_fn(data_item): images, _ = data_item return {input_name: images.numpy()}

Step 2: Initialize NNCF Dataset

calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)

Step 3: Run the quantization pipeline

quantized_model = nncf.quantize(onnx_model, calibration_dataset)

Training-Time Quantization

Here is an example of Accuracy Aware Quantization pipeline where model weights and compression parameters may be fine-tuned to achieve a higher accuracy.

PyTorch

import nncf import nncf.torch import torch from torchvision import datasets, models

Instantiate your uncompressed model

model = models.mobilenet_v2()

Provide validation part of the dataset to collect statistics needed for the compression algorithm

val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()])) dataset_loader = torch.utils.data.DataLoader(val_dataset)

Step 1: Initialize the transformation function

def transform_fn(data_item): images, _ = data_item return images

Step 2: Initialize NNCF Dataset

calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)

Step 3: Run the quantization pipeline

quantized_model = nncf.quantize(model, calibration_dataset)

Now use compressed_model as a usual torch.nn.Module

to fine-tune compression parameters along with the model weights

Save quantization modules and the quantized model parameters

checkpoint = { 'state_dict': model.state_dict(), 'nncf_config': nncf.torch.get_config(model), ... # the rest of the user-defined objects to save } torch.save(checkpoint, path_to_checkpoint)

...

Load quantization modules and the quantized model parameters

resuming_checkpoint = torch.load(path_to_checkpoint) nncf_config = resuming_checkpoint['nncf_config'] state_dict = resuming_checkpoint['state_dict']

quantized_model = nncf.torch.load_from_config(model, nncf_config, example_input) model.load_state_dict(state_dict)

... the rest of the usual PyTorch-powered training pipeline

Training-Time Compression

Here is an example of Accuracy Aware RB Sparsification pipeline where model weights and compression parameters may be fine-tuned to achieve a higher accuracy.

PyTorch

import torch import nncf.torch # Important - must be imported before any other external package that depends on torch

from nncf import NNCFConfig from nncf.torch import create_compressed_model, register_default_init_args

Instantiate your uncompressed model

from torchvision.models.resnet import resnet50 model = resnet50()

Load a configuration file to specify compression

nncf_config = NNCFConfig.from_json("resnet50_imagenet_rb_sparsity.json")

Provide data loaders for compression algorithm initialization, if necessary

import torchvision.datasets as datasets representative_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()])) init_loader = torch.utils.data.DataLoader(representative_dataset) nncf_config = register_default_init_args(nncf_config, init_loader)

Apply the specified compression algorithms to the model

compression_ctrl, compressed_model = create_compressed_model(model, nncf_config)

Now use compressed_model as a usual torch.nn.Module

to fine-tune compression parameters along with the model weights

... the rest of the usual PyTorch-powered training pipeline

Export to ONNX or .pth when done fine-tuning

compression_ctrl.export_model("compressed_model.onnx") torch.save(compressed_model.state_dict(), "compressed_model.pth")

NOTE (PyTorch): Due to the way NNCF works within the PyTorch backend, import nncf must be done before any other import of torch in your package or in third-party packages that your code utilizes. Otherwise, the compression may be applied incompletely.

Tensorflow

import tensorflow as tf

from nncf import NNCFConfig from nncf.tensorflow import create_compressed_model, register_default_init_args

Instantiate your uncompressed model

from tensorflow.keras.applications import ResNet50 model = ResNet50()

Load a configuration file to specify compression

nncf_config = NNCFConfig.from_json("resnet50_imagenet_rb_sparsity.json")

Provide dataset for compression algorithm initialization

representative_dataset = tf.data.Dataset.list_files("/path/*.jpeg") nncf_config = register_default_init_args(nncf_config, representative_dataset, batch_size=1)

Apply the specified compression algorithms to the model

compression_ctrl, compressed_model = create_compressed_model(model, nncf_config)

Now use compressed_model as a usual Keras model

to fine-tune compression parameters along with the model weights

... the rest of the usual TensorFlow-powered training pipeline

Export to Frozen Graph, TensorFlow SavedModel or .h5 when done fine-tuning

compression_ctrl.export_model("compressed_model.pb", save_format="frozen_graph")

For a more detailed description of NNCF usage in your training code, see this tutorial.

Demos, Tutorials and Samples

For a quicker start with NNCF-powered compression, try sample notebooks and scripts presented below.

Jupyter* Notebook Tutorials and Demos

Ready-to-run Jupyter* notebook tutorials and demos are available to explain and display NNCF compression algorithms for optimizing models for inference with the OpenVINO Toolkit:

Notebook Tutorial Name	Compression Algorithm	Backend	Domain
BERT Quantization	Post-Training Quantization	OpenVINO	NLP
MONAI Segmentation Model Quantization	Post-Training Quantization	OpenVINO	Segmentation
PyTorch Model Quantization	Post-Training Quantization	PyTorch	Image Classification
YOLOv11 Quantization with Accuracy Control	Post-Training Quantization with Accuracy Control	OpenVINO	Speech-to-Text,Object Detection
PyTorch Training-Time Compression	Training-Time Compression	PyTorch	Image Classification
TensorFlow Training-Time Compression	Training-Time Compression	Tensorflow	Image Classification
Joint Pruning, Quantization and Distillation for BERT	Joint Pruning, Quantization and Distillation	OpenVINO	NLP

A list of notebooks demonstrating OpenVINO conversion and inference together with NNCF compression for models from various domains:

Demo Model	Compression Algorithm	Backend	Domain
YOLOv8	Post-Training Quantization	OpenVINO	Object Detection,KeyPoint Detection,Instance Segmentation
EfficientSAM	Post-Training Quantization	OpenVINO	Image Segmentation
Segment Anything Model	Post-Training Quantization	OpenVINO	Image Segmentation
OneFormer	Post-Training Quantization	OpenVINO	Image Segmentation
InstructPix2Pix	Post-Training Quantization	OpenVINO	Image-to-Image
CLIP	Post-Training Quantization	OpenVINO	Image-to-Text
BLIP	Post-Training Quantization	OpenVINO	Image-to-Text
Latent Consistency Model	Post-Training Quantization	OpenVINO	Text-to-Image
ControlNet QR Code Monster	Post-Training Quantization	OpenVINO	Text-to-Image
SDXL-turbo	Post-Training Quantization	OpenVINO	Text-to-Image,Image-to-Image
Distil-Whisper	Post-Training Quantization	OpenVINO	Speech-to-Text
Whisper	Post-Training Quantization	OpenVINO	Speech-to-Text
MMS Speech Recognition	Post-Training Quantization	OpenVINO	Speech-to-Text
Grammar Error Correction	Post-Training Quantization	OpenVINO	NLP, Grammar Correction
LLM Instruction Following	Weight Compression	OpenVINO	NLP, Instruction Following
LLM Chat Bots	Weight Compression	OpenVINO	NLP, Chat Bot

Post-Training Quantization Examples

Compact scripts demonstrating quantization and corresponding inference speed boost:

Training-Time Compression Examples

Examples of full pipelines including compression, training, and inference for classification, detection, and segmentation tasks:

Third-party repository integration

NNCF may be easily integrated into training/evaluation pipelines of third-party repositories.

Used by

• OpenVINO Training Extensions
  NNCF is integrated into OpenVINO Training Extensions as a model optimization backend. You can train, optimize, and export new models based on available model templates as well as run the exported models with OpenVINO.
• HuggingFace Optimum Intel
  NNCF is used as a compression backend within the renowned transformers repository in HuggingFace Optimum Intel.
• Ultralytics
  NNCF is integrated into the Intel OpenVINO export pipeline, enabling quantization for the exported models.
• ExecuTorch
  NNCF is used as primary quantization framework for the ExecuTorch OpenVINO integration.
• torch.compile
  NNCF is used as primary quantization framework for the torch.compile OpenVINO integration.

Installation Guide

For detailed installation instructions, refer to the Installation guide.

NNCF can be installed as a regular PyPI package via pip:

NNCF is also available via conda:

conda install -c conda-forge nncf

System requirements of NNCF correspond to the used backend. System requirements for each backend and the matrix of corresponding versions can be found in installation.md.

NNCF Compressed Model Zoo

List of models and compression results for them can be found at our NNCF Model Zoo page.

Citing

@article{kozlov2020neural, title = {Neural network compression framework for fast model inference}, author = {Kozlov, Alexander and Lazarevich, Ivan and Shamporov, Vasily and Lyalyushkin, Nikolay and Gorbachev, Yury}, journal = {arXiv preprint arXiv:2002.08679}, year = {2020} }

Contributing Guide

Refer to the CONTRIBUTING.md file for guidelines on contributions to the NNCF repository.

Useful links

• Documentation
• Example scripts (model objects available through links in respective README.md files):
  • PyTorch
  • TensorFlow
• FAQ
• Notebooks
• HuggingFace Optimum Intel
• OpenVINO Model Optimization Guide

Telemetry

NNCF as part of the OpenVINO™ toolkit collects anonymous usage data for the purpose of improving OpenVINO™ tools. You can opt-out at any time by running the following command in the Python environment where you have NNCF installed:

opt_in_out --opt_out

More information available on OpenVINO telemetry.