目錄
  • 文件 >
  • 訓練後量化 (PTQ)
快捷方式

訓練後量化 (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.com.tw/tutorials/advanced/cpp_frontend.html#loading-data 和 https://pytorch.com.tw/tutorials/recipes/recipes/loading_data_recipe.html。Torch-TensorRT 使用 Dataloaders 作為通用校準器實現的基礎。因此，您將能夠重用或快速為目標域實現 torch::Dataset，將其放入 DataLoader 中，並建立一個 INT8 校準器，您可以將其提供給 Torch-TensorRT，以便在編譯模組時執行 INT8 校準。

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

這是一個 CIFAR10 的 torch::Dataset 類示例介面

 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，該 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 資料載入器與 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). 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.

文件

訪問 PyTorch 的全面開發者文件

檢視文件

教程

獲取面向初學者和高階開發者的深度教程

檢視教程

資源

查詢開發資源並獲得問題解答

檢視資源

為了分析流量和最佳化您的體驗，我們在此網站上提供 Cookie。透過點選或導航，您同意我們使用 Cookie。作為本網站的當前維護者，Facebook 的 Cookie 政策適用。瞭解更多，包括可用的控制選項：Cookie Policy。