訓練後量化 (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 使用數據加載器作為通用校準器實作的基礎。因此，您將能夠重複使用或快速實作目標領域的 torch::Dataset，將其放入數據加載器中，並建立一個 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 數據集的二進制分佈中讀取數據，並建構兩個張量來保存圖像和標籤。

我們使用數據集的一個子集來進行校準，因為我們不需要完整的數據集來進行有效的校準，而且校準確實需要一些時間，然後定義要應用於數據集中圖像的預處理，並從數據集中建立一個數據加載器，該加載器將對數據進行批次處理

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 數據加載器與 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。

訓練後量化 (PTQ)¶

如何在 C++ 中建立自己的 PTQ 應用程式¶

如何在 Python 中建立自己的 PTQ 應用程式¶

引用¶

文件

教學課程

資源