量化：

收集统计信息

运行模型前向传播：
使用校验数据集运行原始浮点模型，记录各层激活值的统计信息：

• 最大值（max）和最小值（min）：用于确定量化范围。
• 直方图分布：分析激活值的分布形态（如对称或偏态）。
• 均值（mean）和方差（variance）：辅助优化量化参数。

对称量化：

基于数组R中最大值，-max|Xf|是通过R中最大值取负号得到的。
优点： 没有变差Z 计算量更小

非对称量化

优点：
非对称量化在数据中不存在离散点，更倾向于饱和量化。

数据出现离散点

当数据中出现比较远的离散点，如果将离散点当做Rmax或者Rmin计算，那么Scale就会差异比较大，最后会出现饱和区间和不饱和区间。

动态范围的常用计算方法

可以完全划分

不能完全划分：

PTQ和QAT的区别

TensoRT 量化方式总结

隐式量化

隐式量化：精度不可控， 存在某一层那个类型快就选择那一层。

隐式量化通常指的是在不修改或微调模型权重的情况下，直接对模型进行量化。这种方法依赖于校准数据集来计算量化参数（如缩放因子和零点）。由于没有对模型进行进一步的训练或微调，量化后的模型精度可能不如原始浮点模型。

特点
无需微调：不需要对模型进行额外的训练或调整。
简单快速：由于不需要微调，整个过程相对简单且快速。
精度不可控：由于没有通过微调来优化量化参数，精度损失可能会较大，尤其是在某些层中表现尤为明显。
校准过程
隐式量化依赖于校准数据集来确定每个层的最佳量化参数。常见的校准方法包括：

最大最小值法：基于校准数据集中的最大值和最小值来确定缩放因子。
熵校准法：通过最小化量化后的分布与原始分布之间的差异来确定最佳缩放因子。

显式量化

原理
显式量化是指在量化之前或之后对模型进行微调（Fine-tuning），以恢复因量化而损失的精度。这种方法通常使用原始数据进行微调，确保量化后的模型尽可能接近原始浮点模型的性能。

特点
需要微调：需要对量化后的模型进行微调，以优化量化参数并恢复精度。
精度可控：通过微调可以显著减少量化带来的精度损失，使量化后的模型在实际应用中仍能保持较高的准确性。
复杂度较高：相比于隐式量化，显式量化的过程更为复杂，因为需要设计合适的微调策略和训练流程。
微调过程
显式量化通常涉及以下几个步骤：

量化感知训练（Quantization Aware Training, QAT）：在训练过程中模拟量化过程，使得模型能够适应量化带来的变化。
后量化微调（Post-Training Fine-Tuning）：在量化完成后对模型进行微调，以恢复精度。

QAT需要进行微调的主要原因

1. 量化带来误差

量化噪声：将浮点数转换为低精度表示（如INT8）时，不可避免地会产生舍入误差。这种误差会影响模型的输出，特别是在敏感的任务中。
激活值范围变化：量化后的激活值范围可能与原始浮点模型不同，这可能导致某些层的输入分布发生变化，影响模型的表现。

代码演示

显式QAT量化案例代码

import torch
import torchvision
from pytorch_quantization import tensor_quant
from pytorch_quantization import quant_modules
from pytorch_quantization import nn as quant_nn
from pytorch_quantization.nn.modules import _utils as  quant_nn_utils
from pytorch_quantization import calib
from typing import List, Callable, Union, Dict

class disable_quantization:
    def __init__(self, model):
        self.model  = model

    def apply(self, disabled=True):
        for name, module in self.model.named_modules():
            if isinstance(module, quant_nn.TensorQuantizer):
                module._disabled = disabled

    def __enter__(self):
        self.apply(True)

    def __exit__(self, *args, **kwargs):
        self.apply(False)


class enable_quantization:
    def __init__(self, model):
        self.model  = model

    def apply(self, enabled=True):
        for name, module in self.model.named_modules():
            if isinstance(module, quant_nn.TensorQuantizer):
                module._disabled = not enabled

    def __enter__(self):
        self.apply(True)
        return self

    def __exit__(self, *args, **kwargs):
        self.apply(False)


def quantizer_state(module):
    for name, module in module.named_modules():
        if isinstance(module, quant_nn.TensorQuantizer):
            print(name, module)


def transfer_torch_to_quantization(nninstance : torch.nn.Module, quantmodule):

    quant_instance = quantmodule.__new__(quantmodule)
    for k, val in vars(nninstance).items():
        setattr(quant_instance, k, val)

    def __init__(self):
      
        if isinstance(self, quant_nn_utils.QuantInputMixin):
            quant_desc_input = quant_nn_utils.pop_quant_desc_in_kwargs(self.__class__,input_only=True)
            self.init_quantizer(quant_desc_input)

            # Turn on torch_hist to enable higher calibration speeds
            if isinstance(self._input_quantizer._calibrator, calib.HistogramCalibrator):
                self._input_quantizer._calibrator._torch_hist = True
        else:
            quant_desc_input, quant_desc_weight = quant_nn_utils.pop_quant_desc_in_kwargs(self.__class__)
            self.init_quantizer(quant_desc_input, quant_desc_weight)

            # Turn on torch_hist to enable higher calibration speeds
            if isinstance(self._input_quantizer._calibrator, calib.HistogramCalibrator):
                self._input_quantizer._calibrator._torch_hist = True
                self._weight_quantizer._calibrator._torch_hist = True

    __init__(quant_instance)
    return quant_instance


def replace_to_quantization_module(model : torch.nn.Module, ignore_policy : Union[str, List[str], Callable] = None):

    module_dict = {}
    for entry in quant_modules._DEFAULT_QUANT_MAP:
        module = getattr(entry.orig_mod, entry.mod_name)
        module_dict[id(module)] = entry.replace_mod

    def recursive_and_replace_module(module, prefix=""):
        for name in module._modules:
            submodule = module._modules[name]
            path      = name if prefix == "" else prefix + "." + name
            recursive_and_replace_module(submodule, path)

            submodule_id = id(type(submodule))
            if submodule_id in module_dict:  
                module._modules[name] = transfer_torch_to_quantization(submodule, module_dict[submodule_id])

    recursive_and_replace_module(model)


#quant_modules.initialize()
model = torchvision.models.resnet50()
model.cuda()
# disable_quantization(model.conv1).apply()
# quantizer_state(model)
replace_to_quantization_module(model)
inputs = torch.randn(1, 3, 224, 224, device='cuda')
quant_nn.TensorQuantizer.use_fb_fake_quant =True
torch.onnx.export(model, inputs, 'quant_resnet50_replace_to_quantization.onnx',opset_version=13)

总结

1. 定义模型。
2. 对模型插入QDQ。
3. 统计QDQ节点的range 和 scale -----> 标定
4. 做敏感层的分析 ------> 需要知道，那层对精度指标影响比较大？？关闭掉对精度影响大的层。
5. 导出一个带QDQ节点的PTQ模型。
6. 对模型进行Finetune ------> QAT

Nvidia 官网案例

#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import datetime
import inspect
import os
import sys
import time
import argparse
import warnings
import collections

import subprocess

import torch
import torch.utils.data

from collections import namedtuple
from torch import nn

from tqdm import tqdm

import torchvision
from torchvision import transforms
from torch.hub import load_state_dict_from_url

from pytorch_quantization import nn as quant_nn
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor
from pytorch_quantization import quant_modules

import onnxruntime
import numpy as np
import models.classification as models

from prettytable import PrettyTable

# The following path assumes running in nvcr.io/nvidia/pytorch:20.08-py3
sys.path.insert(0, "/opt/pytorch/vision/references/classification/")

# Import functions from torchvision reference
try:
    from train import evaluate, train_one_epoch, load_data, utils
except Exception as e:
    raise ModuleNotFoundError(
        "Add https://github.com/pytorch/vision/blob/master/references/classification/ to PYTHONPATH")


def get_parser():
    """
    Creates an argument parser.
    """
    parser = argparse.ArgumentParser(description='Classification quantization flow script')

    parser.add_argument('--data-dir', '-d', type=str, help='input data folder', required=True)
    parser.add_argument('--model-name', '-m', default='resnet50', help='model name: default resnet50')
    parser.add_argument('--disable-pcq',
                        '-dpcq',
                        action="store_true",
                        help='disable per-channel quantization for weights')
    parser.add_argument('--out-dir', '-o', default='/tmp', help='output folder: default /tmp')
    parser.add_argument('--print-freq', '-pf', type=int, default=20, help='evaluation print frequency: default 20')
    parser.add_argument('--threshold',
                        '-t',
                        type=float,
                        default=-1.0,
                        help='top1 accuracy threshold (less than 0.0 means no comparison): default -1.0')

    parser.add_argument('--fp16', action="store_true", help="Enable FP16 model training, evaluation and export")

    parser.add_argument('--batch-size-train', type=int, default=128, help='batch size for training: default 128')
    parser.add_argument('--batch-size-test', type=int, default=128, help='batch size for testing: default 128')
    parser.add_argument('--batch-size-onnx', type=int, default=1, help='batch size for onnx: default 1')

    parser.add_argument('--seed', type=int, default=12345, help='random seed: default 12345')

    checkpoint = parser.add_mutually_exclusive_group(required=True)
    checkpoint.add_argument('--ckpt-path', default='', type=str, help='path to latest checkpoint (default: none)')
    checkpoint.add_argument('--ckpt-url', default='', type=str, help='url to latest checkpoint (default: none)')
    checkpoint.add_argument('--pretrained', action="store_true")

    parser.add_argument('--num-calib-batch',
                        default=4,
                        type=int,
                        help='Number of batches for calibration. 0 will disable calibration. (default: 4)')
    parser.add_argument('--num-finetune-epochs',
                        default=0,
                        type=int,
                        help='Number of epochs to fine tune. 0 will disable fine tune. (default: 0)')
    parser.add_argument('--calibrator', type=str, choices=["max", "histogram"], default="max")
    parser.add_argument('--percentile', nargs='+', type=float, default=[99.9, 99.99, 99.999, 99.9999])
    parser.add_argument('--sensitivity', action="store_true", help="Build sensitivity profile")
    parser.add_argument('--evaluate-onnx', action="store_true", help="Evaluate exported ONNX")
    parser.add_argument('--evaluate-trt', action="store_true", help="Export and evaluate TRT")

    return parser


def prepare_model(model_name,
                  data_dir,
                  per_channel_quantization,
                  batch_size_train,
                  batch_size_test,
                  batch_size_onnx,
                  calibrator,
                  pretrained=True,
                  ckpt_path=None,
                  ckpt_url=None,
                  fp16=False):
    """
    Prepare the model for the classification flow.
    Arguments:
        model_name: name to use when accessing torchvision model dictionary
        data_dir: directory with train and val subdirs prepared "imagenet style"
        per_channel_quantization: iff true use per channel quantization for weights
                                   note that this isn't currently supported in ONNX-RT/Pytorch
        batch_size_train: batch size to use when training
        batch_size_test: batch size to use when testing in Pytorch
        batch_size_onnx: batch size to use when testing with ONNX-RT
        calibrator: calibration type to use (max/histogram)

        pretrained: if true a pretrained model will be loaded from torchvision
        ckpt_path: path to load a model checkpoint from, if not pretrained
        ckpt_url: url to download a model checkpoint from, if not pretrained and no path was given
        * at least one of {pretrained, path, url} must be valid

    The method returns a the following list:
        [
            Model object,
            data loader for training,
            data loader for Pytorch testing,
            data loader for onnx testing
        ]
    """
    # Use 'spawn' to avoid CUDA reinitialization with forked subprocess
    torch.multiprocessing.set_start_method('spawn')

    ## Initialize quantization, model and data loaders
    if per_channel_quantization:
        quant_desc_input = QuantDescriptor(calib_method=calibrator)
        quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
        quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
    else:
        ## Force per tensor quantization for onnx runtime
        quant_desc_input = QuantDescriptor(calib_method=calibrator, axis=None)
        quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
        quant_nn.QuantConvTranspose2d.set_default_quant_desc_input(quant_desc_input)
        quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)

        quant_desc_weight = QuantDescriptor(calib_method=calibrator, axis=None)
        quant_nn.QuantConv2d.set_default_quant_desc_weight(quant_desc_weight)
        quant_nn.QuantConvTranspose2d.set_default_quant_desc_weight(quant_desc_weight)
        quant_nn.QuantLinear.set_default_quant_desc_weight(quant_desc_weight)

    if model_name in models.__dict__:
        model = models.__dict__[model_name](pretrained=pretrained, quantize=True)
    else:
        quant_modules.initialize()
        model = torchvision.models.__dict__[model_name](pretrained=pretrained)
        quant_modules.deactivate()

    if not pretrained:
        if ckpt_path:
            checkpoint = torch.load(ckpt_path)
        else:
            checkpoint = load_state_dict_from_url(ckpt_url)
        if 'state_dict' in checkpoint.keys():
            checkpoint = checkpoint['state_dict']
        elif 'model' in checkpoint.keys():
            checkpoint = checkpoint['model']
        model.load_state_dict(checkpoint)
    model.eval()
    model.cuda()

    if fp16:
        model = model.half()

    ## Prepare the data loaders
    traindir = os.path.join(data_dir, 'train')
    valdir = os.path.join(data_dir, 'val')
    _args = collections.namedtuple("mock_args", [
        "model", "distributed", "cache_dataset", "val_resize_size", "val_crop_size", "train_crop_size", "interpolation",
        "ra_magnitude", "augmix_severity", "weights", "backend", "use_v2"
    ])
    dataset, dataset_test, train_sampler, test_sampler = load_data(
        traindir, valdir,
        _args(model=model_name,
              distributed=False,
              cache_dataset=False,
              val_resize_size=256,
              val_crop_size=224,
              train_crop_size=224,
              interpolation="bilinear",
              ra_magnitude=9,
              augmix_severity=3,
              weights=None,
              backend="pil",
              use_v2=False))

    data_loader_train = torch.utils.data.DataLoader(dataset,
                                                    batch_size=batch_size_train,
                                                    sampler=train_sampler,
                                                    num_workers=4,
                                                    pin_memory=True)

    data_loader_test = torch.utils.data.DataLoader(dataset_test,
                                                   batch_size=batch_size_test,
                                                   sampler=test_sampler,
                                                   num_workers=4,
                                                   pin_memory=True)

    data_loader_onnx = torch.utils.data.DataLoader(dataset_test,
                                                   batch_size=batch_size_onnx,
                                                   sampler=test_sampler,
                                                   num_workers=4,
                                                   pin_memory=True)

    return model, data_loader_train, data_loader_test, data_loader_onnx


def main(cmdline_args):
    parser = get_parser()
    args = parser.parse_args(cmdline_args)
    print(parser.description)
    print(args)

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    ## Prepare the pretrained model and data loaders
    model, data_loader_train, data_loader_test, data_loader_onnx = prepare_model(
        args.model_name, args.data_dir, not args.disable_pcq, args.batch_size_train, args.batch_size_test,
        args.batch_size_onnx, args.calibrator, args.pretrained, args.ckpt_path, args.ckpt_url, args.fp16)

    ## Initial accuracy evaluation
    CrossEntropy = nn.CrossEntropyLoss()

    # nn.CrossEntropyLoss expects float inputs
    def criterion(output, target):
        return CrossEntropy(output.float(), target)

    with torch.no_grad():
        print('Initial evaluation:')
        top1_initial = evaluate(model, criterion, data_loader_test, device="cuda", print_freq=args.print_freq)

    ## Calibrate the model
    with torch.no_grad():
        calibrate_model(model=model,
                        model_name=args.model_name,
                        data_loader=data_loader_train,
                        num_calib_batch=args.num_calib_batch,
                        calibrator=args.calibrator,
                        hist_percentile=args.percentile,
                        out_dir=args.out_dir)

    ## Evaluate after calibration
    if args.num_calib_batch > 0:
        with torch.no_grad():
            print('Calibration evaluation:')
            top1_calibrated = evaluate(model, criterion, data_loader_test, device="cuda", print_freq=args.print_freq)
    else:
        top1_calibrated = -1.0

    ## Build sensitivy profile
    if args.sensitivity:
        build_sensitivity_profile(model, criterion, data_loader_test)

    ## Finetune the model
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.num_finetune_epochs)
    for epoch in range(args.num_finetune_epochs):
        # Training a single epch
        if "print_freq" in inspect.signature(train_one_epoch).parameters:
            train_one_epoch(model, criterion, optimizer, data_loader_train, "cuda", 0, 100)
        else:
            _args = collections.namedtuple("mock_args",
                                           ["print_freq", "clip_grad_norm", "model_ema_steps", "lr_warmup_epochs"])
            train_one_epoch(model, criterion, optimizer, data_loader_train, "cuda", 0,
                            _args(print_freq=100, clip_grad_norm=None, model_ema_steps=32, lr_warmup_epochs=0))
        lr_scheduler.step()

    if args.num_finetune_epochs > 0:
        ## Evaluate after finetuning
        with torch.no_grad():
            print('Finetune evaluation:')
            top1_finetuned = evaluate(model, criterion, data_loader_test, device="cuda")
    else:
        top1_finetuned = -1.0

    ## Export to ONNX
    onnx_filename = args.out_dir + '/' + args.model_name + ".onnx"
    top1_onnx = -1.0
    if args.evaluate_onnx and export_onnx(model, onnx_filename, args.batch_size_onnx, not args.disable_pcq):
        ## Validate ONNX and evaluate
        top1_onnx = evaluate_onnx(onnx_filename, data_loader_onnx, criterion, args.print_freq)

    trt_filename = args.out_dir + '/' + args.model_name + ".trt"
    top1_trt = -1.0
    if args.evaluate_trt and export_trt(model, trt_filename, args.batch_size_onnx, args.fp16):
        ## Validate TRT and evaluate
        top1_trt = evaluate_trt(trt_filename, data_loader_onnx, criterion, args.print_freq)

    ## Print summary
    print("Accuracy summary:")
    table = PrettyTable(['Stage', 'Top1'])
    table.align['Stage'] = "l"
    table.add_row(['Initial', "{:.2f}".format(top1_initial)])
    table.add_row(['Calibrated', "{:.2f}".format(top1_calibrated)])
    table.add_row(['Finetuned', "{:.2f}".format(top1_finetuned)])
    table.add_row(['ONNX', "{:.2f}".format(top1_onnx)])
    if args.evaluate_trt:
        table.add_row(['TRT', "{:.2f}".format(top1_trt)])
    print(table)

    ## Compare results
    if args.threshold >= 0.0:
        if args.evaluate_onnx and top1_onnx < 0.0:
            print("Failed to export/evaluate ONNX!")
            return 1
        if args.evaluate_trt and top1_trt < 0.0:
            print("Failed to export/evaluate TRT!")
            return 1
        if args.num_finetune_epochs > 0:
            if top1_finetuned >= (top1_onnx - args.threshold):
                print("Accuracy threshold was met!")
            else:
                print("Accuracy threshold was missed!")
                return 1

            if args.evaluate_trt and top1_finetuned >= (top1_trt - args.threshold):
                print("TRT Accuracy threshold was met!")
            elif args.evaluate_trt:
                print("TRT Accuracy threshold was missed!")
                return 1

    return 0


def evaluate_onnx(onnx_filename, data_loader, criterion, print_freq):
    """Evaluate accuracy on the given ONNX file using the provided data loader and criterion.
       The method returns the average top-1 accuracy on the given dataset.
    """
    print("Loading ONNX file: ", onnx_filename)
    ort_session = onnxruntime.InferenceSession(onnx_filename, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
    with torch.no_grad():
        metric_logger = utils.MetricLogger(delimiter="  ")
        header = 'Test:'
        with torch.no_grad():
            for image, target in metric_logger.log_every(data_loader, print_freq, header):
                image = image.to("cpu", non_blocking=True)
                image_data = np.array(image)
                input_data = image_data

                # run the data through onnx runtime instead of torch model
                input_name = ort_session.get_inputs()[0].name
                raw_result = ort_session.run([], {input_name: input_data})
                output = torch.tensor((raw_result[0])).float()

                loss = criterion(output, target)
                acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
                batch_size = image.shape[0]
                metric_logger.update(loss=loss.item())
                metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
                metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
        # gather the stats from all processes
        metric_logger.synchronize_between_processes()

        print('  ONNXRuntime: Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}'.format(top1=metric_logger.acc1,
                                                                                              top5=metric_logger.acc5))
        return metric_logger.acc1.global_avg


def evaluate_trt(trt_filename, data_loader, criterion, print_freq):
    print("Loading TRT file: ", trt_filename)

    import pycuda.driver as cuda
    try:
        import pycuda.autoprimaryctx
    except ModuleNotFoundError:
        import pycuda.autoinit

    import tensorrt as trt

    TRT_LOGGER = trt.Logger()

    TRT_tensor = namedtuple('TRT_tensor', ['binding_idx', 'shape', 'dtype', 'device_memory', 'host_memory'])

    def load_engine(engine_file_path):
        assert os.path.exists(engine_file_path)
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())

    def setup_context(engine):
        return engine.create_execution_context()

    def allocate_buffers(engine, context):

        # Allocate host and device buffers
        bindings = []
        inputs = {}
        outputs = {}
        for binding_idx in range(engine.num_bindings):
            binding = engine.get_tensor_name(binding_idx)
            shape = tuple(context.get_tensor_shape(binding))

            size = trt.volume(context.get_tensor_shape(binding))
            dtype = np.dtype(trt.nptype(engine.get_tensor_dtype(binding)))

            device_memory = cuda.mem_alloc(size * dtype.itemsize)
            bindings.append(int(device_memory))

            if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
                inputs[binding] = TRT_tensor(binding_idx, shape, dtype, device_memory, None)
            else:
                host_memory = cuda.pagelocked_empty(size, dtype)
                outputs[binding] = TRT_tensor(binding_idx, shape, dtype, device_memory, host_memory)

        stream = cuda.Stream()
        return bindings, inputs, outputs, stream

    def infer(batch, context, bindings, inputs, outputs, stream):

        # Transfer input data to the GPU.
        for name, trt_in_t in inputs.items():
            buffer = np.ascontiguousarray(batch[name])
            cuda.memcpy_htod_async(trt_in_t.device_memory, buffer, stream)

        # Run inference
        context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)

        # Transfer predictions back from the GPU.
        for _, trt_out_t in outputs.items():
            cuda.memcpy_dtoh_async(trt_out_t.host_memory, trt_out_t.device_memory, stream)

        # Synchronize the stream
        stream.synchronize()

        return {k: torch.tensor(v.host_memory).reshape(v.shape) for k, v in outputs.items()}

    engine = load_engine(trt_filename)
    context = setup_context(engine)
    bindings, inputs, outputs, stream = allocate_buffers(engine, context)

    with torch.no_grad():
        metric_logger = utils.MetricLogger(delimiter="  ")
        header = 'Test:'
        with torch.no_grad():
            for image, target in metric_logger.log_every(data_loader, print_freq, header):
                image = image.to("cpu", non_blocking=True)
                image_data = np.array(image)

                output = infer({"input": image_data}, context, bindings, inputs, outputs, stream)["output"].float()

                loss = criterion(output, target)
                acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
                batch_size = image.shape[0]
                metric_logger.update(loss=loss.item())
                metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
                metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
        # gather the stats from all processes
        metric_logger.synchronize_between_processes()

        print('  TRTRuntime: Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}'.format(top1=metric_logger.acc1,
                                                                                             top5=metric_logger.acc5))
        return metric_logger.acc1.global_avg


def _export_onnx(model, dummy_input, onnx_filename, opset_version):
    try:
        if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
            torch.onnx.export(model,
                              dummy_input,
                              onnx_filename,
                              verbose=False,
                              input_names=["input"],
                              output_names=["output"],
                              opset_version=opset_version,
                              enable_onnx_checker=False,
                              do_constant_folding=True)
        else:
            torch.onnx.export(model,
                              dummy_input,
                              onnx_filename,
                              verbose=False,
                              input_names=["input"],
                              output_names=["output"],
                              opset_version=opset_version,
                              do_constant_folding=True)
    except ValueError:
        print("Failed to export to ONNX")
        return False

    return True


def export_onnx(model, onnx_filename, batch_onnx, per_channel_quantization):
    model.eval()

    if per_channel_quantization:
        opset_version = 13
    else:
        opset_version = 12

    # Export ONNX for multiple batch sizes
    print("Creating ONNX file: " + onnx_filename)
    dummy_input = torch.randn(batch_onnx, 3, 224, 224, device='cuda')  #TODO: switch input dims by model
    return _export_onnx(model, dummy_input, onnx_filename, opset_version)


def export_trt(model, trt_filename, batch_trt, fp16=False):
    model.eval()

    # Export TRT for multiple batch sizes
    print("Creating TRT file: " + trt_filename)
    dummy_input = torch.randn(batch_trt, 3, 224, 224, device='cuda')  #TODO: switch input dims by model

    OPSET = 17
    onnx_filename = trt_filename.replace(".trt", ".onnx")

    if not _export_onnx(model, dummy_input, onnx_filename, OPSET):
        return False

    trt_cmd = f"trtexec --onnx={onnx_filename} --saveEngine={trt_filename} --int8"

    if fp16:
        trt_cmd += " --fp16"

    print(trt_cmd)
    try:
        trt_stdout = subprocess.check_output(trt_cmd, shell=True).decode("utf-8")
    except:
        print("Failed to export to TRT")
        return False

    print(trt_stdout)
    return 'PASSED' in trt_stdout


def calibrate_model(model, model_name, data_loader, num_calib_batch, calibrator, hist_percentile, out_dir):
    """
        Feed data to the network and calibrate.
        Arguments:
            model: classification model
            model_name: name to use when creating state files
            data_loader: calibration data set
            num_calib_batch: amount of calibration passes to perform
            calibrator: type of calibration to use (max/histogram)
            hist_percentile: percentiles to be used for historgram calibration
            out_dir: dir to save state files in
    """

    if num_calib_batch > 0:
        print("Calibrating model")
        with torch.no_grad():
            collect_stats(model, data_loader, num_calib_batch)

        if not calibrator == "histogram":
            compute_amax(model, method="max")
            calib_output = os.path.join(out_dir, F"{model_name}-max-{num_calib_batch*data_loader.batch_size}.pth")
            torch.save(model.state_dict(), calib_output)
        else:
            for percentile in hist_percentile:
                print(F"{percentile} percentile calibration")
                compute_amax(model, method="percentile")
                calib_output = os.path.join(
                    out_dir, F"{model_name}-percentile-{percentile}-{num_calib_batch*data_loader.batch_size}.pth")
                torch.save(model.state_dict(), calib_output)

            for method in ["mse", "entropy"]:
                print(F"{method} calibration")
                compute_amax(model, method=method)
                calib_output = os.path.join(out_dir,
                                            F"{model_name}-{method}-{num_calib_batch*data_loader.batch_size}.pth")
                torch.save(model.state_dict(), calib_output)


def collect_stats(model, data_loader, num_batches):
    """Feed data to the network and collect statistics"""
    # Enable calibrators
    for name, module in model.named_modules():
        if isinstance(module, quant_nn.TensorQuantizer):
            if module._calibrator is not None:
                module.disable_quant()
                module.enable_calib()
            else:
                module.disable()

    # Feed data to the network for collecting stats
    for i, (image, _) in tqdm(enumerate(data_loader), total=num_batches):
        model(image.cuda())
        if i >= num_batches:
            break

    # Disable calibrators
    for name, module in model.named_modules():
        if isinstance(module, quant_nn.TensorQuantizer):
            if module._calibrator is not None:
                module.enable_quant()
                module.disable_calib()
            else:
                module.enable()


def compute_amax(model, **kwargs):
    # Load calib result
    for name, module in model.named_modules():
        if isinstance(module, quant_nn.TensorQuantizer):
            if module._calibrator is not None:
                if isinstance(module._calibrator, calib.MaxCalibrator):
                    module.load_calib_amax()
                else:
                    module.load_calib_amax(**kwargs)
            print(F"{name:40}: {module}")
    model.cuda()


def build_sensitivity_profile(model, criterion, data_loader_test):
    quant_layer_names = []
    for name, module in model.named_modules():
        if name.endswith("_quantizer"):
            module.disable()
            layer_name = name.replace("._input_quantizer", "").replace("._weight_quantizer", "")
            if layer_name not in quant_layer_names:
                quant_layer_names.append(layer_name)
    for i, quant_layer in enumerate(quant_layer_names):
        print("Enable", quant_layer)
        for name, module in model.named_modules():
            if name.endswith("_quantizer") and quant_layer in name:
                module.enable()
                print(F"{name:40}: {module}")
        with torch.no_grad():
            evaluate(model, criterion, data_loader_test, device="cuda")
        for name, module in model.named_modules():
            if name.endswith("_quantizer") and quant_layer in name:
                module.disable()
                print(F"{name:40}: {module}")


if __name__ == '__main__':
    res = main(sys.argv[1:])
    exit(res)