Neural Network Compression Framework (NNCF) provides a suite of post-training and training-time algorithms for neural networks inference optimization in OpenVINO™ with minimal accuracy drop.
NNCF is designed to work with models from PyTorch, 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 at 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. The models created using Sequential or Keras Functional API are only 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, and PyTorch, TensorFlow and ONNX are also supported.
OpenVINO
import nncf
import openvino.runtime as ov
import torch
from torchvision import datasets, transforms
model = ov.Core().read_model("/model_path")
val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()]))
dataset_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1)
def transform_fn(data_item):
images, _ = data_item
return images
calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)
quantized_model = nncf.quantize(model, calibration_dataset)
PyTorch
import nncf
import torch
from torchvision import datasets, models
model = models.mobilenet_v2()
val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()]))
dataset_loader = torch.utils.data.DataLoader(val_dataset)
def transform_fn(data_item):
images, _ = data_item
return images
calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)
quantized_model = nncf.quantize(model, calibration_dataset)
NOTE In case the Post-Training Quantization algorithm could not reach quality requirements you can fine-tune the quantized pytorch model. Example of the Quantization-Aware training pipeline for a pytorch model could be found here.
TensorFlow
import nncf
import tensorflow as tf
import tensorflow_datasets as tfds
model = tf.keras.applications.MobileNetV2()
val_dataset = tfds.load("/path", split="validation",
shuffle_files=False, as_supervised=True)
def transform_fn(data_item):
images, _ = data_item
return images
calibration_dataset = nncf.Dataset(val_dataset, transform_fn)
quantized_model = nncf.quantize(model, calibration_dataset)
ONNX
import onnx
import nncf
import torch
from torchvision import datasets
onnx_model = onnx.load_model("/model_path")
val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()]))
dataset_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1)
input_name = onnx_model.graph.input[0].name
def transform_fn(data_item):
images, _ = data_item
return {input_name: images.numpy()}
calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)
quantized_model = nncf.quantize(onnx_model, calibration_dataset)
Training-Time Quantization
Below 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
model = models.mobilenet_v2()
val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()]))
dataset_loader = torch.utils.data.DataLoader(val_dataset)
def transform_fn(data_item):
images, _ = data_item
return images
calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)
quantized_model = nncf.quantize(model, calibration_dataset)
checkpoint = {
'state_dict': model.state_dict(),
'nncf_config': model.nncf.get_config(),
...
}
torch.save(checkpoint, path_to_checkpoint)
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)
Training-Time Compression
Below 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
from nncf import NNCFConfig
from nncf.torch import create_compressed_model, register_default_init_args
from torchvision.models.resnet import resnet50
model = resnet50()
nncf_config = NNCFConfig.from_json("resnet50_imagenet_rb_sparsity.json")
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)
compression_ctrl, compressed_model = create_compressed_model(model, nncf_config)
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
from tensorflow.keras.applications import ResNet50
model = ResNet50()
nncf_config = NNCFConfig.from_json("resnet50_imagenet_rb_sparsity.json")
representative_dataset = tf.data.Dataset.list_files("/path/*.jpeg")
nncf_config = register_default_init_args(nncf_config, representative_dataset, batch_size=1)
compression_ctrl, compressed_model = create_compressed_model(model, nncf_config)
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
A collection of ready-to-run Jupyter* notebooks tutorials and demos are available to explain and display NNCF compression algorithms for optimizing models for inference with the OpenVINO Toolkit.
Below is a list of notebooks demonstrating OpenVINO conversion and inference together with NNCF compression for models from various domains.
Post-Training Quantization Examples
Compact scripts demonstrating quantization and corresponding inference speed boost:
Training-Time Compression Examples
These examples provide full pipelines including compression, training and inference for classification, detection and segmentation tasks.
Third-party repository integration
NNCF may be straightforwardly integrated into training/evaluation pipelines of third-party repositories.
Used by
-
OpenVINO Training Extensions
NNCF is integrated into OpenVINO Training Extensions as model optimization backend. So you can train, optimize and export new models based on the available model templates as well as run 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 please refer to the Installation page.
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
- Ubuntu* 18.04 or later (64-bit)
- Python* 3.8 or later
- Supported frameworks:
- PyTorch* >=2.2, <2.4
- TensorFlow* >=2.8.4, <=2.12.1
- ONNX* ==1.16.0
- OpenVINO* >=2022.3.0
This repository is tested on Python* 3.8.10, PyTorch* 2.3.0 (NVidia CUDA* Toolkit 12.1) and TensorFlow* 2.12.1 (NVidia CUDA* Toolkit 11.8).
NNCF Compressed Model Zoo
List of models and compression results for them can be found at our 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 is available at https://docs.openvino.ai/latest/openvino_docs_telemetry_information.html.