nncf

nncf

Neural Network Compression Framework for enhanced OpenVINO™ inference

Stars: 962

Visit
 screenshot

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. It is designed to work with models from PyTorch, TorchFX, TensorFlow, ONNX, and OpenVINO™. NNCF offers samples demonstrating compression algorithms for various use cases and models, with the ability to add different compression algorithms easily. It supports GPU-accelerated layers, distributed training, and seamless combination of pruning, sparsity, and quantization algorithms. NNCF allows exporting compressed models to ONNX or TensorFlow formats for use with OpenVINO™ toolkit, and supports Accuracy-Aware model training pipelines via Adaptive Compression Level Training and Early Exit Training.

README:

Neural Network Compression Framework (NNCF)

Key FeaturesInstallationDocumentationUsageTutorials and SamplesThird-party integrationModel Zoo

GitHub Release Website Apache License Version 2.0 PyPI Downloads

Python Backends OS

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

Compression algorithm OpenVINO PyTorch TorchFX TensorFlow ONNX
Post-Training Quantization Supported Supported Experimental Supported Supported
Weights Compression Supported Supported Experimental Not supported Not supported
Activation Sparsity Not supported Experimental Not supported Not supported Not supported

Training-Time Compression Algorithms

Compression algorithm PyTorch TensorFlow
Quantization Aware Training Supported Supported
Mixed-Precision Quantization Supported Not supported
Sparsity Supported Supported
Filter pruning Supported Supported
Movement pruning Experimental Not supported
  • 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.
  • Seamless combination of pruning, sparsity, and quantization algorithms. Please refer to optimum-intel for examples of joint (movement) pruning, quantization, and distillation (JPQD), end-to-end from NNCF optimization to compressed OpenVINO IR.
  • 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.runtime 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
from nncf.torch import disable_patching

# 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)
with nncf.torch.disable_patching():
    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 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': model.nncf.get_config(),
    ... # 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
Colab
Post-Training Quantization OpenVINO NLP
MONAI Segmentation Model Quantization
Binder
Post-Training Quantization OpenVINO Segmentation
PyTorch Model Quantization Post-Training Quantization PyTorch Image Classification
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
Colab
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
Colab
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:

Example Name Compression Algorithm Backend Domain
OpenVINO MobileNetV2 Post-Training Quantization OpenVINO Image Classification
OpenVINO YOLOv8 Post-Training Quantization OpenVINO Object Detection
OpenVINO YOLOv8 QwAС Post-Training Quantization with Accuracy Control OpenVINO Object Detection
OpenVINO Anomaly Classification Post-Training Quantization with Accuracy Control OpenVINO Anomaly Classification
PyTorch MobileNetV2 Post-Training Quantization PyTorch Image Classification
PyTorch SSD Post-Training Quantization PyTorch Object Detection
TorchFX Resnet18 Post-Training Quantization TorchFX Image Classification
TensorFlow MobileNetV2 Post-Training Quantization TensorFlow Image Classification
ONNX MobileNetV2 Post-Training Quantization ONNX Image Classification

Training-Time Compression Examples

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

Example Name Compression Algorithm Backend Domain
PyTorch Image Classification Training-Time Compression PyTorch Image Classification
PyTorch Object Detection Training-Time Compression PyTorch Object Detection
PyTorch Semantic Segmentation Training-Time Compression PyTorch Semantic Segmentation
TensorFlow Image Classification Training-Time Compression TensorFlow Image Classification
TensorFlow Object Detection Training-Time Compression TensorFlow Object Detection
TensorFlow Instance Segmentation Training-Time Compression TensorFlow Instance Segmentation

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.

Installation Guide

For detailed installation instructions, refer to the Installation guide.

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

pip install nncf

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 NNCF 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

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.

For Tasks:

Click tags to check more tools for each tasks

For Jobs:

Alternative AI tools for nncf

Similar Open Source Tools

For similar tasks

For similar jobs