NNI 官方示例教程全景：从 HPO 超参数调优、NAS 架构搜索到模型压缩的完整实战导航

人工智能AutoML机器学习深度学习模型压缩特征工程【免费下载链接】nniAn open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning.项目地址https://gitcode.com/gh_mirrors/nn/nni点击查看免费下载NNINeural Network Intelligence是微软开源的一站式 AutoML 工具包覆盖机器学习生命周期的超参数调优HPO、神经架构搜索NAS、模型压缩剪枝与量化与特征工程。本文以仓库中 docs/source/examples.rst 这一示例索引页为骨架系统梳理其收录的 10 个官方教程它们是理解 NNI 三大核心能力最直接、可复现的入门素材每个教程都提供可下载的 Python 脚本与 Jupyter Notebook。读完本文你将掌握每个示例解决的问题、核心 API 调用与完整运行命令并能直接在仓库的examples/tutorials目录中定位、运行和改造这些示例。一、示例总览三大主题、十个官方教程examples.rst以卡片形式cardlinkitem指令组织教程每个卡片包含标题、描述、跳转链接与所属主题标签。汇总如下主题教程一句话说明可运行源码HPOHPO Quickstart with PyTorch用 HPO 调优 PyTorch FashionMNIST 模型examples/tutorials/hpo_quickstart_pytorch/main.pyHPOHPO Quickstart with TensorFlow用 HPO 调优 TensorFlow MNIST 模型examples/tutorials/hpo_quickstart_tensorflow/main.pyHPOHPO using command line tool用nnictl命令行工具运行 HPO 实验examples/tutorials/hpo_nnictl/nnictl.rstNASHello, NAS!面向初学者的 MNIST 架构搜索教程multi-trial NASexamples/tutorials/hello_nas.pyNASUse NAS Benchmarks as Datasets从 NNI 预处理好的基准数据库中查询流行 NAS 基准数据examples/tutorials/nasbench_as_dataset.pyCompressionGet Started with Model Pruning on MNIST用剪枝压缩模型入门examples/tutorials/pruning_quick_start.pyCompressionGet Started with Model Quantization on MNIST用量化压缩模型入门examples/tutorials/quantization_quick_start.pyCompressionSpeedup Model with Mask剪枝后用 mask 把模型真正变小变快examples/tutorials/pruning_speedup.pyCompressionSpeedup Model with Calibration Config量化后用校准配置把模型真正加速examples/tutorials/quantization_speedup.pyCompressionPruning Bert on Task MNLI端到端演示 NNI 剪枝 Transformer 并展示真实加速examples/tutorials/new_pruning_bert_glue.pyCompressionQuantize Bert on Task MNLI端到端演示 NNI 模拟量化 Transformer 的过程examples/tutorials/quantization_bert_glue.py教程对应的文档页.rst与.ipynb位于 docs/source/tutorials可运行的源码则统一维护在 examples/tutorials。下文按三大主题逐一展开。二、HPO超参数调优快速开始HPO 系列包含三个教程PyTorch 快速开始、TensorFlow 快速开始以及使用命令行工具nnictl的等价版本。三者共享相同的核心流程共 4 步改造模型、定义搜索空间、配置实验、运行实验。2.1 改造模型三行 API 完成自动调优接入HPO 的评估对象是一个独立的模型脚本称为trial code它会被并发执行多次因此必须放在独立文件中。以 examples/tutorials/hpo_quickstart_pytorch/model.py 为例在普通 PyTorch 模型基础上只增加三个 NNI API 调用用nni.get_next_parameter()获取本次要评估的超参数组合用nni.report_intermediate_result()按 epoch 上报中间精度指标用nni.report_final_result()上报最终精度。这三行 API 是 NNI 所有 trial 脚本的通用约定调优器依赖上报结果决定下一步采样方向。2.2 定义搜索空间choice / loguniform / uniform示例准备调优三个超参数features全连接层维度、lr学习率、momentum动量。其先验知识被编码为如下搜索空间定义见 docs/source/tutorials/hpo_quickstart_pytorch/main.rstsearch_space { features: {_type: choice, _value: [128, 256, 512, 1024]}, lr: {_type: loguniform, _value: [0.0001, 0.1]}, momentum: {_type: uniform, _value: [0, 1]}, }其中features是choice从离散候选集中选择lr是loguniform指数分布的对数均匀采样适合跨数量级的学习率momentum是uniform区间内均匀采样。这些类型名源自numpy.random的分布命名。完整的搜索空间规范可参考 docs/source/hpo/search_space.rst。2.3 用 Python API 配置并运行实验NNI 用experiment统一管理 HPO 全过程配置对象定义了如何训练模型与如何探索搜索空间。本示例使用local模式即所有 trial 在本机训练from nni.experiment import Experiment experiment Experiment(local)随后依次配置四块内容# 1) trial 代码每个超参数组合的一次评估称为一个 trial experiment.config.trial_command python model.py experiment.config.trial_code_directory . # 2) 搜索空间 experiment.config.search_space search_space # 3) 调优算法使用 TPE tuner目标是最大化精度 experiment.config.tuner.name TPE experiment.config.tuner.class_args[optimize_mode] maximize # 4) 运行规模总共评估 10 组超参数并发 2 组 experiment.config.max_trial_number 10 experiment.config.trial_concurrency 2需要注意的细节trial_code_directory为相对路径时相对于当前工作目录main.py所在目录若想在别处运行可显式设为Path(__file__).parent。在未安装 Conda 的 Linux 系统上可能需要把python model.py改为python3 model.py。可设置max_experiment_duration 1h限制运行时长若max_trial_number与max_experiment_duration都未设置实验将一直运行直到手动停止。文档明确提示示例中的max_trial_number 10是为了快速演示在默认配置下TPE tuner 需要 20 个 trial 完成热身真实场景应设更大值。启动实验并访问 Web 门户experiment.run(8080) # 浏览器打开 http://localhost:8080 查看状态实验结束后可调用experiment.stop()停止Python 退出时会自动调用可省略若还想查看结果可用experiment.view()重启 Web 门户。完整 Python API 参考见 docs/source/reference/experiment.rst。2.4 用 nnictl 命令行运行同一个实验examples/tutorials/hpo_nnictl/nnictl.rst 演示了与 2.3 完全等价的效果区别在于用命令行工具 YAML 配置文件代替纯 Python 代码。其完整config.yaml如下search_space: features: _type: choice _value: [ 128, 256, 512, 1024 ] lr: _type: loguniform _value: [ 0.0001, 0.1 ] momentum: _type: uniform _value: [ 0, 1 ] trial_command: python model.py trial_code_directory: . trial_concurrency: 2 max_trial_number: 10 tuner: name: TPE class_args: optimize_mode: maximize training_service: platform: local一个关键差异YAML 配置中相对路径的解析规则与 Python API 不同——YAML 里的trial_code_directory相对的是配置文件所在目录因此config.yaml与model.py必须放在同一目录。创建与停止实验$ nnictl create --config config.yaml --port 8080 $ nnictl stop p43ny6ew # 实验 ID 见 create 输出也可用 nnictl stop --all2.5 TensorFlow 版本HPO Quickstart with TensorFlow源码见 examples/tutorials/hpo_quickstart_tensorflow/main.py与 PyTorch 版本流程一致同样通过nni.get_next_parameter/report_intermediate_result/report_final_result三行 API 改造模型同样定义features、lr、momentum的搜索空间用nni.experiment.Experiment配置并运行实验仅将模型框架替换为 TensorFlow。从目录结构看它同样提供main.py与model.py两个文件与 PyTorch 版一一对应。三、NAS从Hello, NAS!到基准数据集NAS 系列包含两个教程Hello, NAS!是神经架构搜索的 101 入门基于 NNI 的 NAS 框架Retiarii在 MNIST 上搜索网络结构nasbench_as_dataset则把 NAS 基准NAS-Bench-101/201、NDS当作数据集来查询。3.1 NAS 任务的三个核心组件examples/tutorials/hello_nas.py文档见 docs/source/tutorials/hello_nas.rst指出一个神经架构搜索任务由三部分组成模型搜索空间Model search space定义要探索的一组模型搜索策略Strategy探索该空间的方法模型评估器Model evaluator报告空间中每个模型的性能。文档同时明确Retiarii 目前仅支持 PyTorch且只在PyTorch 1.9 至 1.13上测试过本教程基于 PyTorch 语境展开。3.2 定义模型空间ModelSpace LayerChoice nni.choice模型空间 一个基础模型 对基础模型的可选变异。定义基础模型与普通 PyTorch 模型几乎相同唯一区别是继承nni.nas.nn.pytorch.ModelSpace而非nn.Moduleimport torch import torch.nn as nn import torch.nn.functional as F import nni from nni.nas.nn.pytorch import LayerChoice, ModelSpace, MutableDropout, MutableLinear class Net(ModelSpace): # 应继承 ModelSpace 而非 nn.Module def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 32, 3, 1) self.conv2 nn.Conv2d(32, 64, 3, 1) self.dropout1 nn.Dropout(0.25) self.dropout2 nn.Dropout(0.5) self.fc1 nn.Linear(9216, 128) self.fc2 nn.Linear(128, 10)在此基础上用变异原语把一个具体模型扩展成包含很多模型的模型空间class DepthwiseSeparableConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.depthwise nn.Conv2d(in_ch, in_ch, kernel_size3, groupsin_ch) self.pointwise nn.Conv2d(in_ch, out_ch, kernel_size1) def forward(self, x): return self.pointwise(self.depthwise(x)) class MyModelSpace(ModelSpace): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 32, 3, 1) # LayerChoice每次采样从候选模块列表中选一个 self.conv2 LayerChoice([ nn.Conv2d(32, 64, 3, 1), DepthwiseSeparableConv(32, 64) ], labelconv2) # nni.choice选一个 dropout 率结果作为 MutableXXX 的参数 self.dropout1 MutableDropout(nni.choice(dropout, [0.25, 0.5, 0.75])) self.dropout2 nn.Dropout(0.5) feature nni.choice(feature, [64, 128, 256]) self.fc1 MutableLinear(9216, feature) self.fc2 MutableLinear(feature, 10)这里用到两个变异 APILayerChoice接受候选模块列表每个被采样的模型从中选一个可像普通 PyTorch 模块一样使用nni.choice作为MutableDropout/MutableLinear的参数把 dropout 率与维度变为可搜索的量。打印该空间会看到Categorical([...], label...)形式的可变维度。若内置变异 API 无法表达你的模型空间可参考 docs/source/nas/mutator.rst 自定义 mutator更多 API 说明见 docs/source/nas/construct_space.rst。3.3 探索模型空间搜索策略 模型评估器NNI NAS 支持多种探索策略见 docs/source/nas/exploration_strategy.rst本教程用多试验multi-trial方式即逐个独立评估采样出的模型另一种是 one-shot 权值共享搜索参考 docs/source/nas/tutorials.rst。实例化策略只需一行import nni.nas.strategy as strategy search_strategy strategy.Random() # 不需要去重时传 dedupFalse评估器用于训练并验证每个模型。教程推荐使用FunctionalEvaluator包装自定义的训练评估函数——该函数接收一个模型类并用nni.report_final_result上报该模型的最终分数。示例评估器在 MNIST 上训练 3 个 epoch 并报告验证精度def evaluate_model(model): device torch.device(cuda) if torch.cuda.is_available() else torch.device(cpu) model.to(device) optimizer torch.optim.Adam(model.parameters(), lr1e-3) transf transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) train_loader DataLoader(MNIST(data/mnist, downloadTrue, transformtransf), batch_size64, shuffleTrue) test_loader DataLoader(MNIST(data/mnist, downloadTrue, trainFalse, transformtransf), batch_size64) for epoch in range(3): train_epoch(model, device, train_loader, optimizer, epoch) accuracy test_epoch(model, device, test_loader) nni.report_intermediate_result(accuracy) # 上报中间结果可为 float 或 dict nni.report_final_result(accuracy) # 上报最终结果 from nni.nas.evaluator import FunctionalEvaluator evaluator FunctionalEvaluator(evaluate_model)推荐让evaluate_model除model外不接受额外参数如需额外参数可参考高级教程 docs/source/nas/evaluator.rst。3.4 启动 NAS 实验、可视化与导出三要素齐备后启动一个 NAS 实验from nni.nas.experiment import NasExperiment exp NasExperiment(model_space, evaluator, search_strategy)与 HPO 实验不同NAS 实验会自动生成一份实验配置多数场景开箱即用例如使用多试验策略时默认使用并发度为 1 的 local 训练服务。也可按需定制exp.config.max_trial_number 3 # 最多生成 3 个 trial exp.config.trial_concurrency 1 # 并发运行 1 个 trial exp.config.trial_gpu_number 0 # 不使用 GPU若想用 GPU包括被占用的 GPU如跑着图形界面的卡需设置exp.config.trial_gpu_number 1 exp.config.training_service.use_active_gpu True启动实验exp.run(port8081)可视化方面与普通 HPO 实验一样打开localhost:8081即可查看 Web 门户同时支持用第三方可视化引擎如 Netron查看每个 trial 的网络结构——点击 trial 详情面板的Visualization按钮即可。当前可视化基于 ONNX因此模型需能导出为 ONNX内置评估器如 Classification会自动导出模型文件自定义评估器则需把模型保存到$NNI_OUTPUT_DIR/model.onnximport os from pathlib import Path def evaluate_model_with_visualization(model): if NNI_OUTPUT_DIR in os.environ: dummy_input torch.zeros(1, 3, 32, 32) torch.onnx.export(model, (dummy_input,), Path(os.environ[NNI_OUTPUT_DIR]) / model.onnx) evaluate_model(model)搜索结束后可用export_top_models导出最优模型for model_dict in exp.export_top_models(formatterdict): print(model_dict) # 示例输出{conv2: 1, dropout: 0.25, feature: 256}NAS 实验同样支持 local 之外的多种训练服务参见 docs/source/experiment/training_service/overview.rst。3.5 把 NAS 基准当作数据集NAS-Bench-101 / 201 / NDSexamples/tutorials/nasbench_as_dataset.py 面向研究场景与其逐一从头训练架构不如直接查询基准数据库中已记录好的架构精度。前置条件是把基准数据准备好放在缓存目录默认~/.cache/nni/nasbenchmark数据准备方法见 docs/source/nas/benchmarks.rst。就绪后目录形如[nds-5745c235.db, nasbench201-b2b60732.db, nasbench101-209f5694.db]NAS-Bench-101 查询用包含操作与输入的字典描述架构调用query_nb101_trial_statsimport pprint from nni.nas.benchmark.nasbench101 import query_nb101_trial_stats arch { op1: conv3x3-bn-relu, op2: maxpool3x3, op3: conv3x3-bn-relu, op4: conv3x3-bn-relu, op5: conv1x1-bn-relu, input1: [0], input2: [1], input3: [2], input4: [0], input5: [0, 3, 4], input6: [2, 5] } for t in query_nb101_trial_stats(arch, 108, include_intermediatesTrue): pprint.pprint(t)返回的生成器中每个元素是一次训练结果包含 train/valid/test 精度、训练时间、epoch 数、参数量等字段{config: {arch: {...}, hash: 00005c142e6f48ac74fdcf73e3439874, id: 4, num_epochs: 108, num_vertices: 7}, id: 10, intermediates: [{current_epoch: 54, test_acc: 77.40, train_acc: 82.82, training_time: 883.46, valid_acc: 77.76}, {current_epoch: 108, test_acc: 92.12, train_acc: 100.0, training_time: 1769.13, valid_acc: 92.42}], parameters: 8.55553, test_acc: 92.12, train_acc: 100.0, training_time: 106147.68, valid_acc: 92.42}NAS-Bench-201 查询以0_1: avg_pool_3x3这类源节点_目标节点操作的字典描述架构可指定数据集如cifar100与训练 epoch 数from nni.nas.benchmark.nasbench201 import query_nb201_trial_stats arch { 0_1: avg_pool_3x3, 0_2: conv_1x1, 1_2: skip_connect, 0_3: conv_1x1, 1_3: skip_connect, 2_3: skip_connect } for t in query_nb201_trial_stats(arch, 200, cifar100): pprint.pprint(t)返回记录还包含flops、latency、各评估集的 loss 等。把num_epochs传None并打开include_intermediatesTrue可拿到每个 epoch 的中间结果imagenet16-120数据集下会返回 200 个中间点。NDS 查询NDSNAS 数据集以model_spec可选与cell_spec可选描述架构bot_muls/ds/num_gs/ss/ws分别代表 bottleneck 倍数、深度、组数、步长与宽度传None表示通配from nni.nas.benchmark.nds import query_nds_trial_stats model_spec {bot_muls: [0.0, 0.25, 0.25, 0.25], ds: [1, 16, 1, 4], num_gs: [1, 2, 1, 2], ss: [1, 1, 2, 2], ws: [16, 64, 128, 16]} for t in query_nds_trial_stats(residual_bottleneck, None, None, model_spec, None, cifar10): pprint.pprint(t)NDS 查询同样支持include_intermediatesTrue获取逐 epoch 指标且可组合各类通配条件例如query_nds_trial_stats(None, amoeba, None, None, None, None, None)统计 amoeba 家族全部样本示例输出为 5107 条。这类查询 API 的底层实现位于 nni/nas/benchmark 目录。四、Compression剪枝、量化与加速Compression 系列共 5 个教程覆盖剪枝/量化 → 生成 mask/校准配置 → 真实加速的完整链路。4.1 剪枝快速开始Pruning Quickstartexamples/tutorials/pruning_quick_start.py 先梳理了剪枝 DNN 的三种常见实践预训练模型 → 剪枝 → 微调剪枝后的模型训练中剪枝pruning aware training→ 微调剪枝后的模型剪枝 → 从零训练剪枝后的模型。NNI 通过支撑关键的剪枝阶段覆盖以上全部实践。教程用 MNIST 预训练的小模型演示流程先预训练 3 个 epoch示例输出精度 72% → 93% → 95%再构造config_list指定剪枝目标。以 L1NormPruner 为例config_list [{ op_types: [Linear, Conv2d], exclude_op_names: [fc3], # fc3 被排除不剪 sparse_ratio: 0.5 # 其余 Linear/Conv2d 层稀疏度 50% }] from nni.compression.pruning import L1NormPruner pruner L1NormPruner(model, config_list)剪枝器会把config_list命中的层用PrunerModuleWrapper包装打印模型可见_nni_wrapper结构随后compress()生成各层 mask_, masks pruner.compress() for name, mask in masks.items(): print(name, sparsity : , {:.2}.format(mask[weight].sum() / mask[weight].numel())) # fc1/conv1/conv2/fc2 的 sparsity 均为 0.5fc3 未被剪config_list的完整规范见 docs/source/compression/config_list.rst。关键一步是用ModelSpeedup把带 mask 的模型压实为真正更小的稠密模型pruner.unwrap_model() # 若模型已被包装需先解包 from nni.compression.speedup import ModelSpeedup ModelSpeedup(model, torch.rand(3, 1, 28, 28).to(device), masks).speedup_model()加速后模型结构发生实质性变化示例中conv1: Conv2d(1, 6) → Conv2d(1, 3)fc1: Linear(256, 120) → Linear(128, 60)等且稀疏度高于 50%——因为ModelSpeedup会把 mask 跨层传播。微调时注意加速会替换层结构必须重新初始化优化器再继续训练。4.2 量化快速开始Quantization Quickstartexamples/tutorials/quantization_quick_start.py 指出量化通过减少表示权重或激活所需的比特数来缩小模型、加速推理。NNI 同时支持训练后量化PTQ与量化感知训练QAT算法教程以QATQuantizer为例import nni from nni.compression.quantization import QATQuantizer from nni.compression.utils import TorchEvaluator optimizer nni.trace(SGD)(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) evaluator TorchEvaluator(training_model, optimizer, training_step) config_list [{ op_names: [conv1, conv2, fc1, fc2], target_names: [_input_, weight, _output_], quant_dtype: int8, quant_scheme: affine, granularity: default, }, { op_names: [relu1, relu2, relu3], target_names: [_output_], quant_dtype: int8, quant_scheme: affine, granularity: default, }] quantizer QATQuantizer(model, config_list, evaluator, len(train_dataloader)) real_input next(iter(train_dataloader))[0].to(device) quantizer.track_forward(real_input) # 用真实输入跟踪 forward 以便收集统计 _, calibration_config quantizer.compress(None, max_epochs5) print(calibration_config)量化训练 5 个 epoch 后得到的calibration_config是一个 defaultdict为每个目标张量权重、输入、输出记录scale、zero_point、quant_dtype、quant_scheme、quant_bits及训练时跟踪到的tracked_max/tracked_min。示例中量化前后精度几乎不变0.9906 → 0.9912。config_list参数op_names/target_names/quant_dtype/quant_scheme/granularity的语义详见 docs/source/compression/config_list.rst。4.3 用 mask 加速剪枝模型ModelSpeedup 原理与独立使用examples/tutorials/pruning_speedup.py 先阐明动机剪枝算法通常用权重 mask 模拟真实剪枝mask 能检验特定稀疏度的性能但没有真实加速。NNI 提供ModelSpeedup把模型转换为真正更小的模型。剪枝分两类细粒度剪枝不改权重与输入/输出张量的形状需要稀疏内核才能加速粗粒度剪枝如通道剪枝会改变张量形状无需稀疏内核直接把被剪层替换为更小层即可。由于社区对稀疏内核的支持有限当前ModelSpeedup仅支持粗粒度剪枝的加速细粒度留待未来。其设计与实现分为两步对应实现位于 nni/compression/speedup形状推断shape inference确定所有应被替换的模块——粗粒度 mask 改变了权重或输入/输出张量形状必须检查是否存在其他未剪层因形状变化也需要替换。拓扑信息来自基于torch.fx的 tracer 得到的模型计算图新模块形状由 NNI 自动推断模块替换module replacement为每类模块提供替换函数返回更小的新模块。ModelSpeedup是相对独立的工具可不依赖 pruner 单独使用手动构造 mask 即可import torch from nni_assets.compression.mnist_model import TorchModel, device model TorchModel().to(device) conv1_mask torch.ones_like(model.conv1.weight.data) conv1_mask[0: 3] 0 # 把 conv1 前 3 个输出通道剪掉 masks {conv1: {weight: conv1_mask}} # mask 格式{layer_name: {weight: ..., bias: ...}} from nni.compression.speedup import ModelSpeedup ModelSpeedup(model, torch.rand(10, 1, 28, 28).to(device), masks).speedup_model() print(model) # conv1: Conv2d(1, 6) → Conv2d(1, 3)conv2 输入通道也由 6 变为 3示例在同一台机器上对比了加速前后的推理耗时conv1剪掉一半输出通道后Original Model - Elapsed Time: 2.30sSpeedup Model - Elapsed Time: 0.094s。文档还给出 legacy 剪枝框架下多组示例的 mask 延迟与加速延迟对比slim、fpgm、l1filter、APoZ 剪枝器详见 docs/source/tutorials/pruning_speedup.rst加速后延迟普遍明显低于 mask 延迟。例如 slim pruner 示例V100 GPU输入torch.randn(64, 3, 32, 32)次数Mask 延迟sSpeedup 延迟s10.011970.00510720.020190.00876940.027330.01480980.043100.027441160.077310.05008320.144640.10027同时需注意其限制当前实现支持PyTorch 1.3.1 或更新版本只能替换模块若forward中的函数需要被替换当前实现无法处理变通办法是把函数改造成 PyTorch 模块不支持的自定义模型需自行实现模块替换函数。4.4 用校准配置加速量化模型TensorRT 两种模式examples/tutorials/quantization_speedup.py 说明量化算法通常以模拟方式进行——把张量量化到目标比特再反量化回 float32用 float32 算子模拟低位计算效果因此量化后的模型并没有延迟收益需要专门的加速阶段。NNI 当前的加速流程为1) 把带量化权重与配置的模型转换为 ONNX 格式2) 把 ONNX 模型送入 TensorRT 生成推理引擎。前置条件强烈推荐使用 NVIDIA 提供的 PyTorch Docker 镜像文档已测试nvcr.io/nvidia/pytorch:22.09-py3启动命令如nvidia-docker run -it nvcr.io/nvidia/pytorch:22.09-py3镜像内需安装nni3.0、pytorch_lightning、pycuda。两种加速模式模式 #1利用 TensorRT 的训练后量化with calibration data。此模式下QAT、LSQ 等量化感知训练算法只负责把权重调整得对量化更友好最后一步量化交给 TensorRT 的 PTQ 完成。先用 128 个样本准备校准数据from nni.compression.quantization_speedup.calibrator import Calibrator calib Calibrator(calib_data, data/calib_cache_file.cache, batch_size64) # TensorRT 按 64 的 batch 处理校准数据 from nni.compression.quantization_speedup import ModelSpeedupTensorRT engine ModelSpeedupTensorRT(model, input_shape(64, 3, 224, 224)) engine.compress_with_calibrator(calib)模式 #2把 TensorRT 当作纯加速后端。不使用 TensorRT 内置 PTQ而是把量化比特宽度和张量取值范围喂给 TensorRT即配置trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS。先用PtqQuantizer得到校准配置再据其生成引擎from nni.compression.quantization import PtqQuantizer config_list [{ quant_types: [input, weight, output], quant_bits: {input: 8, weight: 8, output: 8}, quant_dtype: int, quant_scheme: per_tensor_symmetric, op_types: [default] }] quantizer PtqQuantizer(model, config_list, predict_func, True) quantizer.compress() calibration_config quantizer.export_model() # 注意需重新实例化原始模型——校准配置来自 bn folding 之后而 TensorRT 自己做 bn folding model MobileNetV2() # 重新加载 engine ModelSpeedupTensorRT(model, input_shape(64, 3, 224, 224), configcalibration_config) engine.compress()模式 #2 中有一个易踩的坑校准配置是在 bn folding 之后得到的bn folding 会改变模型结构与权重而 TensorRT 会自行完成 bn folding因此必须把原始模型而非折叠后的模型交给 TensorRT。相关实现位于 nni/compression/quantization_speedup。4.5 BERT 端到端示例剪枝 MNLI 与量化 GLUE两个 BERT 教程是 Compression 系列中的进阶端到端示例均基于 HuggingFacetransformers生态。Pruning Bert on Task MNLIexamples/tutorials/new_pruning_bert_glue.py是 NNI v3.0 的 Transformer 剪枝新教程与旧版相比集成了融合压缩剪枝 蒸馏能力、使用更强大稳定的剪枝加速工具并且额外剪掉了整个模型的隐藏维度embedding 层从而大幅缩小模型体积。剪枝过程分三步1) 剪 attention 层2) 剪 feed-forward 层3) 剪 embedding 层。每步先由剪枝器做模拟剪枝生成 mask再做加速——通过稀疏传播发掘局部 mask 带来的全局冗余最终把模型替换为更小的版本。压缩天然搭配蒸馏方法教程同时使用 distiller 帮助恢复精度。工程细节上教程用transformers.Trainer减轻训练/评估逻辑负担但需要用nni.trace包装 Trainer 以跟踪其初始化参数——因为训练感知式剪枝与蒸馏需要重新创建 trainer。Quantize Bert on Task GLUEexamples/tutorials/quantization_bert_glue.py展示了一套 NNI 团队实践过的有效 Transformer 模拟量化流程用户可在此基础上探索更优方案。流程为1) 用 BERT-base-uncased transformers 的 trainer 管道在 GLUE 下游任务上微调经验表明在微调后的模型上量化的最终效果优于直接对预训练模型量化2) 用具体 quantizer 在 GLUE 上量化微调模型教程对比了 QAT、LSQ均为量化感知训练与 PTQ训练后量化三种方法在量化 BERT 上的表现。量化过程中用 int8 量化 BERT encoder 中的 Linear 层。教程默认dev_mode True各训练 1 个 epoch 以便生成文档正式运行需设为False微调与量化各 10 个 epoch。五、在仓库中定位并运行这些示例所有教程均有三种形态可按需选用可运行源码.py位于 examples/tutorials可直接执行。例如python examples/tutorials/pruning_quick_start.py。文档页.rst与Jupyter Notebook.ipynb位于 docs/source/tutorials由 Sphinx-Gallery 从源码自动生成每篇文档末尾附有下载入口。依赖的辅助代码与数据剪枝/加速教程复用了 nni_assets/compression/mnist_model.pyTorchModel、trainer、evaluator、deviceNAS 教程的基准查询底层实现在 nni/nas/benchmark。HPO 实验的另一种等价实现是 YAML 配置 nnictl命令行配置示例见 examples/tutorials/hpo_nnictl。教程涉及的核心 API 参考文档实验管理见 docs/source/reference/experiment.rstnnictl命令见 docs/source/reference/nnictl.rst剪枝器与量化器列表见 docs/source/reference/compressionNAS 空间构建见 docs/source/nas/construct_space.rst。六、结语这 10 个官方示例构成了从入门到进阶的完整学习路径HPO 系列展示了接入自动调优只需三行 API并提供了 Python API 与nnictl命令行两套等价用法NAS 系列演示了用ModelSpace 变异原语构建搜索空间、用策略与评估器驱动搜索以及把 NAS 基准当作数据集进行查询Compression 系列则完整覆盖剪枝/量化 → 生成 mask/校准配置 →ModelSpeedup/TensorRT 真实加速的链路并延伸到 BERT 这类真实 Transformer 模型。每个示例都可以在 examples/tutorials 中找到源码直接运行是上手 NNI 与深入理解其内部实现的最佳起点。赞分享人工智能AutoML机器学习深度学习模型压缩特征工程【免费下载链接】nniAn open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning.项目地址https://gitcode.com/gh_mirrors/nn/nni点击查看免费下载相关推荐NNI 官方教程导航全览从 HPO 调参到 NAS 搜索与模型压缩的实战路径NNI 官方教程导航全览从 HPO 调参到 NAS 搜索与模型压缩的实战路径 本文基于 NNI 开源仓库 docs/source/tutorials/inde人工智能AutoML机器学习深度学习模型压缩特征工程NNI 架构总览从核心概念到端到端调参、NAS 与模型压缩的实现原理NNI 架构总览从核心概念到端到端调参、NAS 与模型压缩的实现原理 NNINeural Network Intelligence是微软开源的一套 Aut人工智能AutoML机器学习深度学习模型压缩特征工程NNI NAS 实战教程从构造模型空间到复现 DARTS 搜索NNI NAS 实战教程从构造模型空间到复现 DARTS 搜索 本教程围绕 NNIRetiariiNAS 框架的两篇官方实战指南展开先通过 MNIST人工智能AutoML机器学习深度学习模型压缩特征工程创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考