首页>编程相关>Ai源码
下载

神经网络压缩框架（NNCF）

主要功能 • 安装 • 文档 • 使用 • 教程和示例 • 第三方集成 • Model Zoo

神经网络压缩框架 (NNCF) 提供了一套训练后和训练时算法，用于优化 OpenVINO™ 中神经网络的推理，同时将精度下降降至最低。

NNCF 旨在与 PyTorch、TorchFX、TensorFlow、ONNX 和 OpenVINO™ 的模型配合使用。

NNCF 提供了演示压缩算法在不同用例和模型中的用法的示例。在 NNCF Model Zoo 页面上查看使用 NNCF 支持的示例可实现的压缩结果。

该框架被组织为一个 Python* 包，可以在独立模式下构建和使用。框架架构是统一的，可以轻松为 PyTorch 和 TensorFlow 深度学习框架添加不同的压缩算法。

主要特点

训练后压缩算法

压缩算法开放VINO火炬火炬FX TensorFlow奥恩克斯
训练后量化支持支持实验性的支持支持
权重压缩支持支持实验性的不支持不支持
激活稀疏性不支持实验性的不支持不支持不支持

训练时间压缩算法

压缩算法火炬TensorFlow
量化意识训练支持支持
混合精度量化支持不支持
稀疏性支持支持
过滤器修剪支持支持
运动修剪实验性的不支持
  • 自动、可配置的模型图转换以获得压缩模型。

    注意：对 TensorFlow 模型的支持有限。仅支持使用 Sequential 或 Keras Function API 创建的模型。

  • 压缩方法的通用接口。
  • GPU 加速层可实现更快的压缩模型微调。
  • 分布式培训支持。
  • 著名第三方存储库 (huggingface-transformers) 的 Git 补丁演示了将 NNCF 集成到自定义训练管道中的过程。
  • 剪枝、稀疏和量化算法的无缝组合。请参阅optimum-intel，了解联合（运动）修剪、量化和蒸馏 (JPQD)、从 NNCF 优化到压缩 OpenVINO IR 的端到端示例。
  • 将 PyTorch 压缩模型导出到 ONNX* 检查点，并将 TensorFlow 压缩模型导出到 SavedModel 或 Frozen Graph 格式，以便与 OpenVINO™ 工具套件一起使用。
  • 通过自适应压缩级别训练和早期退出训练支持准确性感知模型训练管道。

文档

本文档涵盖了有关 NNCF 算法和为 NNCF 做出贡献所需的函数的详细信息。

NNCF 的最新用户文档可在此处获取。

NNCF API 文档可在此处找到。

用法

训练后量化

NNCF PTQ 是应用 8 位量化的最简单方法。要运行该算法，您只需要模型和一个小型（约 300 个样本）校准数据集。

OpenVINO 是运行 PTQ 的首选后端，同时也支持 PyTorch、TensorFlow 和 ONNX。

开放VINO 
 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 )
火炬
 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 )

注意如果训练后量化算法不满足质量要求，您可以微调量化的 pytorch 模型。您可以在此处找到 pytorch 模型的量化感知训练管道的示例。

火炬FX 
 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 )
奥恩克斯
 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 )

训练时间量化

这是准确性感知量化管道的示例，其中可以微调模型权重和压缩参数以实现更高的准确性。

火炬
 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

训练时间压缩

这是一个 Accuracy Aware RB Sparsification pipeline 的示例，其中可以微调模型权重和压缩参数以实现更高的精度。

火炬
 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" )

注意 (PyTorch) ：由于 NNCF 在 PyTorch 后端的工作方式， import nncf必须在您的包或代码使用的第三方包中导入任何其他torch之前完成。否则，压缩可能会不完全。

张量流
 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" )

有关训练代码中 NNCF 用法的更详细说明，请参阅本教程。

演示、教程和示例

为了更快地开始使用 NNCF 支持的压缩，请尝试下面提供的示例笔记本和脚本。

Jupyter* 笔记本教程和演示

可立即运行的 Jupyter* 笔记本教程和演示可解释和展示 NNCF 压缩算法，以优化模型以使用 OpenVINO 工具包进行推理：

笔记本教程名称压缩算法后端领域
BERT 量化
训练后量化开放VINO自然语言处理
MONAI 分割模型量化
训练后量化开放VINO分割
PyTorch 模型量化训练后量化火炬图像分类
具有精度控制的量化具有精度控制的训练后量化开放VINO语音转文本，
物体检测
PyTorch 训练时间压缩训练时间压缩火炬图像分类
TensorFlow 训练时间压缩训练时间压缩张量流图像分类
BERT 的联合剪枝、量化和蒸馏联合剪枝、量化和蒸馏开放VINO自然语言处理

演示 OpenVINO 转换和推理以及针对不同领域模型的 NNCF 压缩的笔记本列表：

演示模型压缩算法后端领域
YOLOv8
训练后量化开放VINO物体检测，
关键点检测，
实例分割
高效SAM训练后量化开放VINO图像分割
分段任何模型训练后量化开放VINO图像分割
前任训练后量化开放VINO图像分割
指导Pix2Pix训练后量化开放VINO图像到图像
夹子训练后量化开放VINO图像到文本
BLIP训练后量化开放VINO图像到文本
潜在一致性模型训练后量化开放VINO文本转图像
ControlNet QR 码怪物训练后量化开放VINO文本转图像
SDXL-涡轮训练后量化开放VINO文本到图像，
图像到图像
蒸馏耳语训练后量化开放VINO语音转文本
耳语
训练后量化开放VINO语音转文本
彩信语音识别训练后量化开放VINO语音转文本
语法错误纠正训练后量化开放VINO NLP、语法纠正
LLM指令遵循重量压缩开放VINO NLP，指令遵循
LLM 聊天机器人重量压缩开放VINO NLP、聊天机器人

训练后量化示例

演示量化和相应推理速度提升的紧凑脚本：

示例名称压缩算法后端领域
OpenVINO MobileNetV2训练后量化开放VINO图像分类
OpenVINO YOLOv8训练后量化开放VINO物体检测
OpenVINO YOLOv8 QwAС具有精度控制的训练后量化开放VINO物体检测
OpenVINO 异常分类具有精度控制的训练后量化开放VINO异常分类
PyTorch MobileNetV2训练后量化火炬图像分类
PyTorch SSD训练后量化火炬物体检测
TorchFX Resnet18训练后量化火炬FX图像分类
TensorFlow MobileNetV2训练后量化TensorFlow图像分类
ONNX MobileNetV2训练后量化奥恩克斯图像分类

训练时间压缩示例

完整管道的示例，包括用于分类、检测和分割任务的压缩、训练和推理：

示例名称压缩算法后端领域
PyTorch 图像分类训练时间压缩火炬图像分类
PyTorch 对象检测训练时间压缩火炬物体检测
PyTorch 语义分割训练时间压缩火炬语义分割
TensorFlow 图像分类训练时间压缩TensorFlow图像分类
TensorFlow 对象检测训练时间压缩TensorFlow物体检测
TensorFlow 实例分割训练时间压缩TensorFlow实例分割

第三方存储库集成

NNCF 可以轻松集成到第三方存储库的培训/评估管道中。

使用者

  • OpenVINO 培训扩展

    NNCF 作为模型优化后端集成到 OpenVINO Training Extensions 中。您可以根据可用的模型模板训练、优化和导出新模型，以及使用 OpenVINO 运行导出的模型。

  • HuggingFace 最佳英特尔

    NNCF 在 HuggingFace Optimum Intel 著名的transformers存储库中用作压缩后端。

安装指南

有关详细的安装说明，请参阅安装指南。

NNCF 可以通过 pip 作为常规 PyPI 包安装： 

pip install nncf

NNCF 也可以通过 conda 获得： 

conda install -c conda-forge nncf

系统要求

  • Ubuntu* 18.04 或更高版本（64 位）
  • Python* 3.9 或更高版本
  • 支持的框架：
    • PyTorch* >=2.4, <2.6
    • TensorFlow* >=2.8.4，<=2.15.1
    • ONNX* ==1.17.0
    • OpenVINO* >=2022.3.0

该存储库已在 Python* 3.10.14、PyTorch* 2.5.0 (NVidia CUDA* Toolkit 12.4) 和 TensorFlow* 2.12.1 (NVidia CUDA* Toolkit 11.8) 上进行测试。

NNCF 压缩 NNCF 模型动物园

模型列表及其压缩结果可以在我们的 NNCF Model Zoo 页面找到。

引用

@ 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 }
}

贡献指南

有关对 NNCF 存储库贡献的指南，请参阅 CONTRIBUTING.md 文件。

有用的链接

  • 文档
  • 示例脚本（可通过相应 README.md 文件中的链接获取模型对象）：
    • 火炬
    • TensorFlow
  • 常问问题
  • 笔记本电脑
  • HuggingFace 最佳英特尔
  • OpenVINO 模型优化指南

遥测

NNCF 作为 OpenVINO™ 工具套件的一部分，收集匿名使用数据，以改进 OpenVINO™ 工具。您可以随时通过在安装了 NNCF 的 Python 环境中运行以下命令来选择退出：

opt_in_out --opt_out

有关 OpenVINO 遥测的更多信息。

展开
附加信息
相关应用
为您推荐
相关资讯 全部