训练后量化 (PTQ)

训练后量化（PTQ）是一种技术，用于减少推理所需的计算资源，同时通过将传统的FP32激活空间映射到减少的INT8空间来保持模型的准确性。TensorRT使用一个校准步骤，该步骤使用目标域的样本数据执行您的模型，并跟踪FP32中的激活，以校准到INT8的映射，从而最小化FP32推理和INT8推理之间的信息损失。

编写TensorRT应用程序的用户需要设置一个校准器类，该类将向TensorRT校准器提供样本数据。通过Torch-TensorRT，我们希望利用PyTorch中的现有基础设施，使实现校准器变得更加容易。

LibTorch 提供了 DataLoader 和 Dataset API，这些 API 简化了预处理和批量输入数据的过程。这些 API 通过 C++ 和 Python 接口暴露，使得最终用户更容易使用。对于 C++ 接口，我们使用 torch::Dataset 和 torch::data::make_data_loader 对象来构建并对数据集进行预处理。Python 接口中的等效功能使用 torch.utils.data.Dataset 和 torch.utils.data.DataLoader。PyTorch 文档的这一部分提供了更多信息 https://pytorch.org/tutorials/advanced/cpp_frontend.html#loading-data 和 https://pytorch.org/tutorials/recipes/recipes/loading_data_recipe.html。Torch-TensorRT 使用 Dataloaders 作为通用校准器实现的基础。因此，您将能够重用或快速实现一个 torch::Dataset 用于您的目标领域，将其放入 DataLoader 中，并创建一个 INT8 校准器，您可以将其提供给 Torch-TensorRT 以在模块编译期间运行 INT8 校准。

如何在C++中创建你自己的PTQ应用程序

这是一个用于CIFAR10的torch::Dataset类的示例接口：

//cpp/ptq/datasets/cifar10.h
#pragma once

#include "torch/data/datasets/base.h"
#include "torch/data/example.h"
#include "torch/types.h"

#include <cstddef>
#include <string>

namespace datasets {
// The CIFAR10 Dataset
class CIFAR10 : public torch::data::datasets::Dataset<CIFAR10> {
public:
    // The mode in which the dataset is loaded
    enum class Mode { kTrain, kTest };

    // Loads CIFAR10 from un-tarred file
    // Dataset can be found https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz
    // Root path should be the directory that contains the content of tarball
    explicit CIFAR10(const std::string& root, Mode mode = Mode::kTrain);

    // Returns the pair at index in the dataset
    torch::data::Example<> get(size_t index) override;

    // The size of the dataset
    c10::optional<size_t> size() const override;

    // The mode the dataset is in
    bool is_train() const noexcept;

    // Returns all images stacked into a single tensor
    const torch::Tensor& images() const;

    // Returns all targets stacked into a single tensor
    const torch::Tensor& targets() const;

    // Trims the dataset to the first n pairs
    CIFAR10&& use_subset(int64_t new_size);


private:
    Mode mode_;
    torch::Tensor images_, targets_;
};
} // namespace datasets

此类的实现从CIFAR10数据集的二进制分发中读取，并构建两个张量，分别保存图像和标签。

我们使用数据集的一个子集进行校准，因为我们不需要完整的数据集来进行有效的校准，而且校准确实需要一些时间，然后定义应用于数据集中图像的预处理，并从数据集中创建一个DataLoader，该DataLoader将对数据进行批处理：

auto calibration_dataset = datasets::CIFAR10(data_dir, datasets::CIFAR10::Mode::kTest)
                                    .use_subset(320)
                                    .map(torch::data::transforms::Normalize<>({0.4914, 0.4822, 0.4465},
                                                                            {0.2023, 0.1994, 0.2010}))
                                    .map(torch::data::transforms::Stack<>());
auto calibration_dataloader = torch::data::make_data_loader(std::move(calibration_dataset),
                                                            torch::data::DataLoaderOptions().batch_size(32)
                                                                                            .workers(2));

接下来，我们使用校准器工厂（位于torch_tensorrt/ptq.h中）从calibration_dataloader创建一个校准器：

#include "torch_tensorrt/ptq.h"
...

auto calibrator = torch_tensorrt::ptq::make_int8_calibrator(std::move(calibration_dataloader), calibration_cache_file, true);

在这里，我们还定义了一个位置来写入校准缓存文件，我们可以使用该文件来重用校准数据，而无需数据集，以及如果缓存文件存在，我们是否应该使用它。还存在一个torch_tensorrt::ptq::make_int8_cache_calibrator工厂，它创建了一个校准器，该校准器仅在您可能在存储有限的机器上（即没有空间存储完整数据集）进行引擎构建或为了拥有更简单的部署应用程序时使用缓存。

校准器工厂创建了一个继承自nvinfer1::IInt8Calibrator虚拟类（默认情况下为nvinfer1::IInt8EntropyCalibrator2）的校准器，该类定义了校准时使用的校准算法。您可以像这样明确选择校准算法：

// MinMax Calibrator is geared more towards NLP tasks
auto calibrator = torch_tensorrt::ptq::make_int8_calibrator<nvinfer1::IInt8MinMaxCalibrator>(std::move(calibration_dataloader), calibration_cache_file, true);

然后，为INT8校准设置模块所需的所有操作就是在torch_tensorrt::CompileSpec结构中设置以下编译设置并编译模块：

std::vector<std::vector<int64_t>> input_shape = {{32, 3, 32, 32}};
/// Configure settings for compilation
auto compile_spec = torch_tensorrt::CompileSpec({input_shape});
/// Set operating precision to INT8
compile_spec.enabled_precisions.insert(torch::kF16);
compile_spec.enabled_precisions.insert(torch::kI8);
/// Use the TensorRT Entropy Calibrator
compile_spec.ptq_calibrator = calibrator;

auto trt_mod = torch_tensorrt::CompileGraph(mod, compile_spec);

如果您已经有一个用于TensorRT的Calibrator实现，您可以直接将ptq_calibrator字段设置为指向您的校准器的指针，它也将正常工作。 从这里开始，执行方式没有太大变化。您仍然可以完全使用LibTorch作为推理的唯一接口。数据在传递到trt_mod.forward时应保持FP32精度。Torch-TensorRT演示中有一个示例应用程序，它带您从在CIFAR10上训练VGG16网络到使用Torch-TensorRT部署在INT8中：https://github.com/pytorch/TensorRT/tree/master/cpp/ptq

如何在Python中创建你自己的PTQ应用程序

Torch-TensorRT Python API 提供了一种简单方便的方法，将 PyTorch 数据加载器与 TensorRT 校准器一起使用。DataLoaderCalibrator 类可以通过提供所需的配置来创建 TensorRT 校准器。以下代码演示了如何使用它的示例

testing_dataset = torchvision.datasets.CIFAR10(
    root="./data",
    train=False,
    download=True,
    transform=transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ]
    ),
)

testing_dataloader = torch.utils.data.DataLoader(
    testing_dataset, batch_size=1, shuffle=False, num_workers=1
)
calibrator = torch_tensorrt.ptq.DataLoaderCalibrator(
    testing_dataloader,
    cache_file="./calibration.cache",
    use_cache=False,
    algo_type=torch_tensorrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
    device=torch.device("cuda:0"),
)

trt_mod = torch_tensorrt.compile(model, inputs=[torch_tensorrt.Input((1, 3, 32, 32))],
                                    enabled_precisions={torch.float, torch.half, torch.int8},
                                    calibrator=calibrator,
                                    device={
                                         "device_type": torch_tensorrt.DeviceType.GPU,
                                         "gpu_id": 0,
                                         "dla_core": 0,
                                         "allow_gpu_fallback": False,
                                         "disable_tf32": False
                                     })

在用户希望使用预先存在的校准缓存文件的情况下，CacheCalibrator 可以在没有任何数据加载器的情况下使用。以下示例演示了如何在 INT8 模式下使用 CacheCalibrator。

calibrator = torch_tensorrt.ptq.CacheCalibrator("./calibration.cache")

trt_mod = torch_tensorrt.compile(model, inputs=[torch_tensorrt.Input([1, 3, 32, 32])],
                                      enabled_precisions={torch.float, torch.half, torch.int8},
                                      calibrator=calibrator)

如果您已经有一个现有的校准器类（直接使用TensorRT API实现），您可以直接将校准器字段设置为您自己的类，这将非常方便。 有关如何使用Torch-TensorRT API在VGG网络上执行PTQ的演示，您可以参考https://github.com/pytorch/TensorRT/blob/master/tests/py/test_ptq_dataloader_calibrator.py 和https://github.com/pytorch/TensorRT/blob/master/tests/py/test_ptq_trt_calibrator.py

引用

Krizhevsky, A., & Hinton, G. (2009). 从微小图像中学习多层特征。

Simonyan, K., & Zisserman, A. (2014). 用于大规模图像识别的非常深的卷积网络。arXiv预印本 arXiv:1409.1556。