訓練後量化 (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.dev.org.tw/tutorials/advanced/cpp_frontend.html#loading-data 和 https://pytorch.dev.org.tw/tutorials/recipes/recipes/loading_data_recipe.html。Torch-TensorRT 使用 Dataloader 作為通用校準器實作的基礎。因此，您可以重複使用或快速實作針對您的目標領域的 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，它將批次處理資料

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。以下範例示範了如何在 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). Learning multiple layers of features from tiny images.

Simonyan, K., & Zisserman, A. (2014). Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556.