AIMET API 文档（7）

- - 1.1.10 自动定量 API
  - 1.1.11 BatchNorm 重新估计 API

1.1.10 自动定量 API

用户指南链接

要了解有关此技术的更多信息，请参阅AutoQuant

示例笔记本链接

有关展示如何使用 PyTorch AutoQuant 的端到端笔记本，请参阅此处。

顶级API

类 aimet_torch.auto_quant_v2.AutoQuant（model、dummy_input、data_loader、eval_callback、param_bw=8、output_bw=8、quant_scheme=<QuantScheme.post_training_tf_enhanced: 2>、rounding_mode=‘nearest’、config_file=None、results_dir=‘/tmp’、cache_id=None、strict_validation=正确）

集成并应用训练后量化技术。

AutoQuant 包括 1）batchnorm 折叠、2）跨层均衡和 3）Adaround。这些技术将以尽力而为的方式应用，直到模型满足 allowed_accuracy_drop 给出的评估目标。

参数

model ( Module) – 要量化的模型。假设模型位于正确的设备上
dummy_input ( Union[ Tensor, Tuple]) – 模型的虚拟输入。假设 dummy_input 在正确的设备上
data_loader ( DataLoader[+T_co]) – 迭代未标记数据集的集合，用于计算编码
eval_callback ( Callable[[ Module], float]) – 计算评估分数的函数
param_bw ( int) – 参数位宽
output_bw ( int) – 输出位宽
quant_scheme ( QuantScheme) – 量化方案
rounding_mode ( str) – 舍入模式
config_file ( Optional[ str]) – 模型量化器的配置文件路径
results_dir ( str) – 保存 PTQ 技术结果的目录
cache_id ( Optional[ str]) – 与缓存结果关联的 ID
strict_validation ( bool) – 标志默认设置为 True。如果为 False，AutoQuant 将继续执行并在可能的情况下在内部处理错误。这可能会产生不理想或不直观的结果。

run_inference( )

创建量化模型并执行推理

返回类型
Tuple[ QuantizationSimModel, float]

退货
QuantizationSimModel，模型精度为浮点数

optimize( allowed_accuracy_drop=0.0 )

集成并应用训练后量化技术。

参数
allowed_accuracy_drop ( float) – 允许的最大精度下降

返回类型
Tuple[ Module, float, str]

退货
元组（最佳模型、评估分数、编码路径）

set_adaround_params（adaround_params ）

设置 Adaround 参数。如果用户未显式调用此方法，AutoQuant 将为 Adaround 使用 data_loader（传递给__init__）。

参数
adaround_params ( AdaroundParameters) – Adaround 参数。

返回类型
None

set_export_params（onnx_export_args=-1，propagate_encodings=None ）

设置 QuantizationSimModel.export 的参数。

参数

onnx_export_args ( OnnxExportApiArgs) – 如果不通过 torchscript 图表提供导出，则可选导出参数，具有 onnx 特定覆盖
propagate_encodings ( Optional[ bool]) – 如果为 True，则中间操作的编码条目（当一个 PyTorch 操作导致多个 ONNX 节点时）将使用与该系列操作的输出张量相同的 BW 和 data_type 进行填充。

返回类型
None

set_model_preparer_params（modules_to_exclude =无，module_classes_to_exclude =无，concrete_args =无）

设置模型准备器的参数。

参数

module_to_exclude ( Optional[ List[ Module]]) – 跟踪时要排除的模块列表。
module_classes_to_exclude ( Optional[ List[ Module]]) – 跟踪时要排除的模块类列表。
crete_args ( Optional[ Dict[ str, Any]]) – 模型准备器的参数。允许您部分专业化您的功能，无论是删除控制流还是数据结构。如果模型具有控制流，torch.fx 将无法跟踪模型。详细检查 torch.fx.symbolic_trace API。

get_quant_scheme_candidates( )

返回量化方案搜索的候选者。期间optimize()，将在其中选出准确率最高的候选者。

返回类型
Tuple[ _QuantSchemePair，……]

退货
量化方案搜索的候选者

set_quant_scheme_candidates（候选人）

设置量化方案搜索的候选者。期间optimize()，将在其中选出准确率最高的候选者。

参数
Candidates ( Tuple[ _QuantSchemePair, …]) – 量化方案搜索的候选者

类 aimet_torch.auto_quant.AutoQuant( allowed_accuracy_drop、unlabeled_dataset_iterable、eval_callback、default_param_bw=8、default_output_bw=8、default_quant_scheme=<QuantScheme.post_training_tf_enhanced: 2>、default_rounding_mode=‘nearest’、default_config_file=None )

警告
auto_quant.AutoQuantauto_quant_v2.AutoQuant 已弃用，并将在后续版本中替换为。

集成并应用训练后量化技术。

AutoQuant 包括 1）batchnorm 折叠、2）跨层均衡和 3）Adaround。这些技术将以尽力而为的方式应用，直到模型满足 allowed_accuracy_drop 给出的评估目标。

参数

allowed_accuracy_drop ( float) – 允许的最大精度下降。
unlabeled_dataset_iterable ( Union[ DataLoader[+T_co], Collection[+T_co]]) – 一个集合（即可使用__len__迭代），迭代用于编码计算的未标记数据集。该迭代产生的值预计能够直接传递给模型。默认情况下，除非self.set_adaround_params另有指定，否则此迭代也将用于 Adaround 。
eval_callback ( Callable[[ Module, Optional[ int]], float]) – 将模型和样本数映射到评估分数的函数。此回调预计会返回一个标量值，表示针对N 个样本评估的模型性能，其中N是作为此回调的第二个参数传递的样本数。注意：如果N为“无”，则预计将根据整个评估数据集来评估模型。
default_param_bw ( int) – 用于量化图层参数的默认位宽 (4-31)。
default_output_bw ( int) – 用于量化层输入和输出的默认位宽 (4-31)。
default_quant_scheme ( QuantScheme) – 量化方案。支持的值为 QuantScheme.post_training_tf 或 QuantScheme.post_training_tf_enhanced。
default_rounding_mode ( str) – 舍入模式。支持的选项是“最近”或“随机”
default_config_file ( Optional[ str]) – 模型量化器的配置文件路径

apply( fp32_model、dummy_input_on_cpu、dummy_input_on_gpu=None、results_dir=‘/tmp’、cache_id=None )

应用训练后量化技术。

参数

fp32_model ( Module) – 应用 PTQ 技术的模型。
dummy_input_on_cpu ( Union[ Tensor, Tuple]) – CPU 内存中模型的虚拟输入。
dummy_input_on_gpu ( Union[ Tensor, Tuple, None]) – GPU 内存中模型的虚拟输入。当且仅当 fp32_model 在 GPU 上时才需要此参数。
results_dir ( str) – 保存结果的目录。
cache_id ( Optional[ str]) – 与 results_dir 组合组成缓存 id 的字符串。如果指定，如果存在相同 results_dir 和 cache_id 下生成的先前 PTQ 结果，则 AutoQuant 将从文件系统加载 PTQ 结果/将 PTQ 结果保存到文件系统，

返回类型
Tuple[ Module, float, str]

退货
元组（最佳模型、评估分数、编码路径前面）。

提高
如果模型位于 GPU 上且未指定 dummy_input_on_gpu，则会出现 ValueError 错误。

set_adaround_params（adaround_params ）

设置 Adaround 参数。如果用户未显式调用此方法，AutoQuant 将为 Adaround使用unlabeled_dataset_iterable（传递给__init__ ）。

参数
adaround_params ( AdaroundParameters) – Adaround 参数。

返回类型
None

set_export_params（onnx_export_args=-1，propagate_encodings=None ）

设置 QuantizationSimModel.export 的参数。

参数

onnx_export_args ( OnnxExportApiArgs) – 如果不通过 torchscript 图表提供导出，则可选导出参数，具有 onnx 特定覆盖
propagate_encodings ( Optional[ bool]) – 如果为 True，则中间操作的编码条目（当一个 PyTorch 操作导致多个 ONNX 节点时）将使用与该系列操作的输出张量相同的 BW 和 data_type 进行填充。

返回类型
None

代码示例

import random
from typing import Optional

import torch
from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
from torchvision import models, datasets, transforms

from aimet_torch.adaround.adaround_weight import AdaroundParameters
from aimet_torch.auto_quant_v2 import AutoQuant

# Step 1. Define constants and helper functions
EVAL_DATASET_SIZE = 5000
CALIBRATION_DATASET_SIZE = 2000
BATCH_SIZE = 100

_subset_samplers = {
    
    }

def _create_sampled_data_loader(dataset, num_samples):
    if num_samples not in _subset_samplers:
        indices = random.sample(range(len(dataset)), num_samples)
        _subset_samplers[num_samples] = SubsetRandomSampler(indices=indices)
    return DataLoader(dataset,
                      sampler=_subset_samplers[num_samples],
                      batch_size=BATCH_SIZE)

# Step 2. Prepare model and dataset
fp32_model = models.resnet18(pretrained=True).eval()

input_shape = (1, 3, 224, 224)
dummy_input = torch.randn(input_shape)

transform = transforms.Compose((
    transforms.ToTensor(),
))
# NOTE: In the actual use cases, a real dataset should provide by the users.
eval_dataset = datasets.FakeData(size=EVAL_DATASET_SIZE,
                                 image_size=input_shape[1:],
                                 num_classes=1000,
                                 transform=transform)

# Step 3. Prepare unlabeled dataset
# NOTE: In the actual use cases, the users should implement this part to serve
#       their own goals if necessary.
class UnlabeledDatasetWrapper(Dataset):
    def __init__(self, dataset):
        self._dataset = dataset

    def __len__(self):
        return len(self._dataset)

    def __getitem__(self, index):
        images, _ = self._dataset[index]
        return images

unlabeled_dataset = UnlabeledDatasetWrapper(eval_dataset)
unlabeled_data_loader = _create_sampled_data_loader(unlabeled_dataset, CALIBRATION_DATASET_SIZE)

# Step 4. Prepare eval callback
# NOTE: In the actual use cases, the users should implement this part to serve
#       their own goals if necessary.
def eval_callback(model: torch.nn.Module, num_samples: Optional[int] = None) -> float:
    if num_samples is None:
        num_samples = len(eval_dataset)

    eval_data_loader = _create_sampled_data_loader(eval_dataset, num_samples)

    num_correct_predictions = 0
    for images, labels in eval_data_loader:
        predictions = torch.argmax(model(images.cuda()), dim=1)
        num_correct_predictions += torch.sum(predictions.cpu() == labels)

    return int(num_correct_predictions) / num_samples

# Step 5. Create AutoQuant object
auto_quant = AutoQuant(fp32_model.cuda(),
                       dummy_input.cuda(),
                       unlabeled_data_loader,
                       eval_callback)

# Step 6. (Optional) Set adaround params
ADAROUND_DATASET_SIZE = 2000
adaround_data_loader = _create_sampled_data_loader(unlabeled_dataset, ADAROUND_DATASET_SIZE)
adaround_params = AdaroundParameters(adaround_data_loader, num_batches=len(adaround_data_loader))
auto_quant.set_adaround_params(adaround_params)

# Step 7. Run AutoQuant
sim, initial_accuracy = auto_quant.run_inference()
model, optimized_accuracy, encoding_path = auto_quant.optimize(allowed_accuracy_drop=0.01)

print(f"- Quantized Accuracy (before optimization): {initial_accuracy:.4f}")
print(f"- Quantized Accuracy (after optimization):  {optimized_accuracy:.4f}")

笔记
要使用auto_quant.AutoQuant（将被弃用），请将以下代码更改应用于步骤 5 和 7。

# Step 5. Create AutoQuant object
auto_quant = AutoQuant(allowed_accuracy_drop=0.01,
                       unlabeled_dataset_iterable=unlabeled_data_loader,
                       eval_callback=eval_callback)

# Step 6. (Optional) Set adaround params
ADAROUND_DATASET_SIZE = 2000
adaround_data_loader = _create_sampled_data_loader(unlabeled_dataset, ADAROUND_DATASET_SIZE)
adaround_params = AdaroundParameters(adaround_data_loader, num_batches=len(adaround_data_loader))
auto_quant.set_adaround_params(adaround_params)

# Step 7. Run AutoQuant
model, accuracy, encoding_path =\
    auto_quant.apply(fp32_model.cuda(),
                     dummy_input_on_cpu=dummy_input.cpu(),
                     dummy_input_on_gpu=dummy_input.cuda())

print(f"- Quantized Accuracy (after optimization):  {optimized_accuracy:.4f}")

1.1.11 BatchNorm 重新估计 API

示例笔记本链接
有关展示如何使用 PyTorch 量化感知训练和 BatchNorm 重新估计的端到端笔记本，请参阅此处

介绍
Batch Norm (BN) Re-estimation 在执行 QAT 后重新估计 BN 层的统计数据。使用重新估计的统计数据，BN 层被折叠到前面的 Conv 层和 Linear 层中

顶级 API
用于 BatchNorm 重新估计的 API

aimet_torch.bn_reestimation.reestimate_bn_stats（模型、数据加载器、num_batches=100、forward_fn=None ）

重新估计 BatchNorm 统计数据（运行均值和方差）。

参数
model ( Module) – 重新估计 BN 统计数据的模型。

dataloader ( DataLoader[+T_co]) – 训练数据集。

num_batches ( int) – 用于重新估计的批次数。

forward_fn ( Optional[ Callable[[ Module, Any], Any]]) – 可选适配器函数，在给定模型和数据加载器生成的输入批次的情况下执行前向传递。

返回类型
Handle

退货
撤消 BN 重新估计对 handle.remove() 的影响的句柄。

用于 BatchNorm 折叠缩放的 API

aimet_torch.batch_norm_fold.fold_all_batch_norms_to_scale（模拟）
将模型中的所有batch_norm层折叠到相应转换层的量化尺度参数中

参数
sim ( QuantizationSimModel) – QuantizationSimModel

返回类型
List[ Tuple[ QcQuantizeWrapper，QcQuantizeWrapper]]

退货
层对列表 [(Conv/Linear, 折叠的 BN 层)]

代码示例 - BN 重新估计

** 步骤 1. 加载模型**

对于此示例，我们将从 torchvision 加载预训练的 ResNet18 模型。

def load_fp32_model():

    import torchvision
    from torchvision.models import resnet18
    from aimet_torch.model_preparer import prepare_model

    use_cuda = torch.cuda.is_available()
    if use_cuda:
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    model = resnet18(pretrained=True).to(device)
    model = prepare_model(model)

    return model, use_cuda

步骤 2. 创建启用范围学习和每通道量化的 QuantSim

有关使用范围学习 QuantScheme 创建 QuantSim 的示例
有关如何启用每通道量化

步骤 3. 执行 QAT

  # User action required
    # The following line of code is an example of how to use an example ImageNetPipeline's train function.
    # Replace the following line with your own pipeline's  train function.
    ImageNetPipeline.train(sim.model, epochs=1, learning_rate=5e-7, learning_rate_schedule=[5, 10], use_cuda=use_cuda)

步骤 4a. 执行 BatchNorm 重新估计

 from aimet_torch.bn_reestimation import reestimate_bn_stats

    # User action required
    # The following line of code is an example of how to use the ImageNet data's training data loader.
    # Replace the following line with your own dataset's training data loader.
    train_loader = ImageNetDataPipeline.get_train_dataloader()

    reestimate_bn_stats(quant_sim.model, train_loader, forward_fn=forward_fn)

步骤 4 b. 执行 BatchNorm Fold 以缩放

 from aimet_torch.batch_norm_fold import fold_all_batch_norms_to_scale

    fold_all_batch_norms_to_scale(quant_sim)

步骤 5. 导出模型和编码并在目标上进行测试