PyTorch神经网络压缩框架

PyTorch神经网络压缩框架

Neural Network Compression Framework (NNCF)

This repository contains a PyTorch*-based framework and samples for neural networks compression.

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

The samples demonstrate the usage of compression algorithms for three different use cases on public models and datasets: Image Classification, Object Detection and Semantic Segmentation. Compression results achievable with the NNCF-powered samples are reported here.

Key Features

Support of various compression algorithms, applied during a model fine-tuning process to achieve best compression parameters and accuracy: QuantizationBinarizationSparsityFilter pruning Automatic, configurable model graph transformation to obtain the compressed model. The source model is wrapped by the custom class and additional compression-specific layers are inserted in the graph.Common interface for compression methodsGPU-accelerated layers for faster compressed model fine-tuningDistributed training supportConfiguration file examples for each supported compression algorithm.Git patches for prominent third-party repositories (mmdetection, huggingface-transformers) demonstrating the process of integrating NNCF into custom training pipelinesExporting compressed models to ONNX* checkpoints ready for usage with OpenVINO™ toolkit.

Usage

The NNCF is organized as a regular Python package that can be imported in your target training pipeline script. The basic workflow is loading a jsON configuration script containing NNCF-specific parameters determining the compression to be applied to your model, and then passing your model along with the configuration script to the nncf.create_compressed_model function. This function returns a wrapped model ready for compression fine-tuning, and handle to the object allowing you to control the compression during the training process:

import nncffrom nncf import create_compressed_model, Config as NNCFConfig# Instantiate your uncompressed modelfrom torchvision.models.resnet import resnet50model = resnet50()# Load a configuration file to specify compressionnncf_config = NNCFConfig.from_json("resnet50_int8.json")# Provide data loaders for compression algorithm initialization, if necessarynncf_config = register_default_init_args(nncf_config, loss_criterion, train_loader)# Apply the specified compression algorithms to the modelcomp_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-tuningcomp_ctrl.export_model("compressed_model.onnx")torch.save(compressed_model.state_dict(), "compressed_model.pth")

For a more detailed description of NNCF usage in your training code, see Usage.md. For in-depth examples of NNCF integration, browse the sample scripts code, or the example patches to third-party repositories.

For more details about the framework architecture, refer to the NNCFArchitecture.md.

Model Compression Samples

For a quicker start with NNCF-powered compression, you can also try the sample scripts, each of which provides a basic training pipeline for classification, semantic segmentation and object detection neural network training correspondingly.

To run the samples please refer to the corresponding tutorials:

Image Classification sampleObject Detection sampleSemantic Segmentation sample

Third-party repository integration

NNCF may be straightforwardly integrated into training/evaluation pipelines of third-party repositories. See third_party_integration for examples of code modifications (Git patches and base commit IDs are provided) that are necessary to integrate NNCF into select repositories.

System requirements

Ubuntu* 16.04 or later (64-bit)Python* 3.6 or laterNVidia CUDA* Toolkit 10.2 or laterPyTorch* 1.5 or later.

Installation

We suggest to install or use the package in the Python virtual environment.

As a package built from checked-out repository:

Install the following system dependencies:

sudo apt-get install python3-dev

Install the package and its dependencies by running the following in the repository root directory:

For CPU & GPU-powered execution: python setup.py installFor CPU-only installation python setup.py install --cpu-only

As a Docker image

Use one of the Dockerfiles in the docker directory to build an image with an environment already set up and ready for running NNCF sample scripts.

Contributing

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

NNCF compression results

Achieved using sample scripts and NNCF configuration files provided with this repository. See README.md files for sample scripts for links to exact configuration files and final PyTorch checkpoints.

Legal Information

[*] Other names and brands may be claimed as the property of others.

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。