训练后量化 (PTQ)¶
训练后量化 (PTQ) 是一种技术，通过将传统的 FP32 激活空间映射到缩减的 INT8 空间，从而减少推理所需的计算资源，同时保持模型的准确性。TensorRT 使用校准步骤，该步骤使用目标域中的样本数据执行您的模型，并跟踪 FP32 中的激活以校准映射到 INT8，从而最大程度地减少 FP32 推理和 INT8 推理之间的信息损失。
编写 TensorRT 应用程序的用户需要设置一个校准器类，该类将向 TensorRT 校准器提供样本数据。使用 Torch-TensorRT，我们希望利用 PyTorch 中现有的基础设施来简化校准器的实现。
LibTorch 提供了 DataLoader 和 Dataset API，简化了输入数据的预处理和批处理。这些 API 通过 C++ 和 Python 接口公开，使用户更容易使用。对于 C++ 接口，我们使用 torch::Dataset 和 torch::data::make_data_loader 对象来构建和对数据集执行预处理。Python 接口中的等效功能使用 torch.utils.data.Dataset 和 torch.utils.data.DataLoader。PyTorch 文档的这一部分提供了更多信息 https://pytorch.ac.cn/tutorials/advanced/cpp_frontend.html#loading-data 和 https://pytorch.ac.cn/tutorials/recipes/recipes/loading_data_recipe.html。Torch-TensorRT 使用 Dataloader 作为通用校准器实现的基础。因此，您将能够重用或快速实现目标域的 torch::Dataset，将其放入 DataLoader 中并创建一个 INT8 校准器，您可以将其提供给 Torch-TensorRT 以在编译模块期间运行 INT8 校准。
如何在 C++ 中创建自己的 PTQ 应用程序¶
这是一个 torch::Dataset 类的 CIFAR10 接口示例
1//cpp/ptq/datasets/cifar10.h
2#pragma once
3
4#include "torch/data/datasets/base.h"
5#include "torch/data/example.h"
6#include "torch/types.h"
7
8#include <cstddef>
9#include <string>
10
11namespace datasets {
12// The CIFAR10 Dataset
13class CIFAR10 : public torch::data::datasets::Dataset<CIFAR10> {
14public:
15 // The mode in which the dataset is loaded
16 enum class Mode { kTrain, kTest };
17
18 // Loads CIFAR10 from un-tarred file
19 // Dataset can be found https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz
20 // Root path should be the directory that contains the content of tarball
21 explicit CIFAR10(const std::string& root, Mode mode = Mode::kTrain);
22
23 // Returns the pair at index in the dataset
24 torch::data::Example<> get(size_t index) override;
25
26 // The size of the dataset
27 c10::optional<size_t> size() const override;
28
29 // The mode the dataset is in
30 bool is_train() const noexcept;
31
32 // Returns all images stacked into a single tensor
33 const torch::Tensor& images() const;
34
35 // Returns all targets stacked into a single tensor
36 const torch::Tensor& targets() const;
37
38 // Trims the dataset to the first n pairs
39 CIFAR10&& use_subset(int64_t new_size);
40
41
42private:
43 Mode mode_;
44 torch::Tensor images_, targets_;
45};
46} // namespace datasets
此类的实现从 CIFAR10 数据集的二进制分布中读取，并构建两个张量，分别保存图像和标签。
我们使用数据集的一个子集进行校准，因为我们不需要整个数据集来进行有效的校准，并且校准确实需要一些时间，然后定义要应用于数据集中的图像的预处理，并从数据集中创建 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 校准器实现，您可以直接使用指向校准器的指针设置 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 dataloader 与 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 而无需任何 dataloader。以下示例演示了如何使用 CacheCalibrator 在 INT8 模式下使用。
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 直接实现），您可以直接将校准器字段设置为您的类，这非常方便。有关如何在 VGG 网络上使用 Torch-TensorRT API 执行 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。
