TensorRT部署系列 | 如何將模型從 PyTorch 轉(zhuǎn)換為 TensorRT 并加速推理？

機(jī)器學(xué)習(xí)工程師的生活包括長(zhǎng)時(shí)間的挫折和片刻的歡樂(lè)！

首先，努力讓你的模型在你的訓(xùn)練數(shù)據(jù)上產(chǎn)生好的結(jié)果。您可視化您的訓(xùn)練數(shù)據(jù)，清理它，然后再次訓(xùn)練。您閱讀了機(jī)器學(xué)習(xí)中的偏差方差權(quán)衡（bias variance tradeoff）以系統(tǒng)地處理訓(xùn)練過(guò)程。

有一天，你的 PyTorch 模型經(jīng)過(guò)完美訓(xùn)練，可以投入生產(chǎn)了。

那是純粹的快樂(lè)！

您對(duì)準(zhǔn)確性感到自豪，您在項(xiàng)目跟蹤器中將您的任務(wù)標(biāo)記為已完成，并通知您的 CTO 模型已準(zhǔn)備就緒。

她不贊成地?fù)u搖頭，告訴你這個(gè)模型還沒(méi)有在生產(chǎn)環(huán)境上準(zhǔn)備好！訓(xùn)練模型是不夠的。您需要修改模型，使其在運(yùn)行（也稱(chēng)為推理）時(shí)高效。

你不知道如何進(jìn)行。您好心的 CTO 告訴您在 http://LearnOpenCV.com 上閱讀這篇關(guān)于 TensorRT 的帖子。因此，在這里您將對(duì)另一種學(xué)習(xí)體驗(yàn)感到高興。

在本文中，如果您已經(jīng)在PyTorch中訓(xùn)練了網(wǎng)絡(luò)，您將學(xué)習(xí)如何快速輕松地使用「TensorRT」進(jìn)行部署。

我們將使用以下步驟。

1. 使用 PyTorch 訓(xùn)練模型
2. 將模型轉(zhuǎn)換為 ONNX 格式
3. 使用 NVIDIA TensorRT 進(jìn)行推理

在本教程中，我們僅使用預(yù)訓(xùn)練模型并跳過(guò)步驟 1?，F(xiàn)在，讓我們了解什么是 ONNX 和 TensorRT。

1、什么是 ONNX？

有許多用于訓(xùn)練深度學(xué)習(xí)模型的框架。最受歡迎的是 Tensorflow 和 PyTorch。但是，由 Tensorflow 訓(xùn)練的模型不能與 PyTorch 一起使用，反之亦然。

ONNX 代表開(kāi)放神經(jīng)網(wǎng)絡(luò)交換。它是一種用于表示機(jī)器學(xué)習(xí)模型的開(kāi)放格式。

您可以在您選擇的任何框架中訓(xùn)練您的模型，然后將其轉(zhuǎn)換為 ONNX 格式。

擁有通用格式的巨大好處是，在運(yùn)行時(shí)加載模型的軟件或硬件只需要與 ONNX 兼容。

ONNX 之于機(jī)器學(xué)習(xí)模型就像 JPEG 之于圖像或 MPEG 之于視頻。

2、什么是 TensorRT？

NVIDIA 的 TensorRT 是一個(gè)用于高性能深度學(xué)習(xí)推理的 SDK。

它提供 API 來(lái)對(duì)預(yù)訓(xùn)練模型進(jìn)行推理，并為您的平臺(tái)生成優(yōu)化的運(yùn)行時(shí)引擎。

有多種方法可以實(shí)現(xiàn)這種優(yōu)化。例如，TensorRT 使我們能夠使用 INT8（8 位整數(shù)）或 FP16（16 位浮點(diǎn)數(shù)）運(yùn)算，而不是通常的 FP32。這種精度的降低可以顯著加快推理速度，但精度會(huì)略有下降。

其他類(lèi)型的優(yōu)化包括通過(guò)重用內(nèi)存、融合層和張量、根據(jù)硬件選擇合適的數(shù)據(jù)層等來(lái)最大限度地減少 GPU 內(nèi)存占用。

3、TensorRT 的環(huán)境設(shè)置

要重現(xiàn)本文中提到的實(shí)驗(yàn)，您需要NVIDIA顯卡。任何比 Maxwell（算力5.0）更新的架構(gòu)都可以。您可以在此處的表格中找到您的 GPU 計(jì)算能力：https://developer.nvidia.com/cuda-gpus#compute。不要忘記安裝合適的驅(qū)動(dòng)程序。

3.1 安裝 PyTorch、ONNX 和 OpenCV

安裝「Python 3.6」或更高版本并運(yùn)行

python3 -m pip install -r requirements.txt

requirements.txt內(nèi)容：

torch==1.2.0
torchvision==0.4.0
albumentations==0.4.5
onnx==1.4.1
opencv-python==4.2.0.34

代碼在指定版本上進(jìn)行了測(cè)試。但如果您已經(jīng)安裝了其中一些組件，則可以嘗試在其他版本上啟動(dòng)它。

3.2 安裝 TensorRT

1. 按照官方說(shuō)明下載并安裝NVIDIA CUDA 10.0或更高版本：https://developer.nvidia.com/cuda-10.0-download-archive
2. 下載并提取適用于您的 CUDA 版本的CuDNN庫(kù)（需要登錄）：https://developer.nvidia.com/rdp/cudnn-download
3. 下載并提取適用于您的 CUDA 版本的 NVIDIA TensorRT庫(kù)（需要登錄）：https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html。所需的最低版本為 6.0.1.5。請(qǐng)按照您系統(tǒng)的安裝指南(https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html)進(jìn)行操作，不要忘記安裝Python 的部分
4. 將 CUDA、TensorRT、CuDNN 庫(kù)的絕對(duì)路徑添加到環(huán)境變量PATH或LD_LIBRARY_PATH
5. 安裝PyCUDA(https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-pycuda)

我們現(xiàn)在準(zhǔn)備好進(jìn)行我們的實(shí)驗(yàn)。

4、如何將 PyTorch 模型轉(zhuǎn)換為 TensorRT

讓我們回顧一下將 PyTorch 模型轉(zhuǎn)換為 TensorRT 所需的步驟。

1. 使用 PyTorch 加載并啟動(dòng)預(yù)訓(xùn)練模型

首先，讓我們?cè)?PyTorch 上使用預(yù)訓(xùn)練網(wǎng)絡(luò)實(shí)現(xiàn)一個(gè)簡(jiǎn)單的分類(lèi)。例如，我們將采用Resnet50，但您可以選擇任何您想要的。您可以在此處找到有關(guān)如何使用 PyTorch 的更多信息和解釋?zhuān)? PyTorch for Beginners: Image Classification using Pre-trained models

from torchvision import models
model = models.resnet50(pretrained=True)

下一個(gè)重要步驟：「預(yù)處理」輸入圖像。我們需要知道在訓(xùn)練期間進(jìn)行了哪些轉(zhuǎn)換以在推理的時(shí)候復(fù)制它們。我們推薦以下模塊用于預(yù)處理步驟：「albumentations」和「cv2」 (OpenCV)。

該模型在大小為 224×224 的圖像上進(jìn)行訓(xùn)練。然后將輸入數(shù)據(jù)歸一化（將像素值除以 255，減去平均值并除以標(biāo)準(zhǔn)差）。

import cv2
import torch
from albumentations import Resize, Compose
from albumentations.pytorch.transforms import  ToTensor
from albumentations.augmentations.transforms import Normalize
 
def preprocess_image(img_path):
    # transformations for the input data
    transforms = Compose([
        Resize(224, 224, interpolation=cv2.INTER_NEAREST),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ToTensor(),
    ])
     
    # read input image
    input_img = cv2.imread(img_path)
    # do transformations
    input_data = transforms(image=input_img)["image"]

準(zhǔn)備批次以傳遞到網(wǎng)絡(luò)。在我們的案例中，批處理中只有一張圖像。請(qǐng)注意，我們將輸入數(shù)據(jù)上傳到 GPU 以更快地執(zhí)行程序，使我們與 TensorRT 的比較更加公平。

    batch_data = torch.unsqueeze(input_data, 0)
    return batch_data
 
input = preprocess_image("turkish_coffee.jpg").cuda()

現(xiàn)在我們可以進(jìn)行推理了。不要忘記將模型切換到評(píng)估模式并將其也復(fù)制到 GPU。結(jié)果，我們將得到對(duì)象屬于哪個(gè)類(lèi)的概率 tensor[1, 1000]。

model.eval()
model.cuda()
output = model(input)

為了獲得人類(lèi)可讀的結(jié)果，我們需要后處理步驟。分類(lèi)標(biāo)簽可以在imagenet_classes.txt中找到。計(jì)算Softmax以獲得每個(gè)類(lèi)別的百分比并打印網(wǎng)絡(luò)預(yù)測(cè)的最高類(lèi)別。

def postprocess(output_data):
    # get class names
    with open("imagenet_classes.txt") as f:
        classes = [line.strip() for line in f.readlines()]
    # calculate human-readable value by softmax
    confidences = torch.nn.functional.softmax(output_data, dim=1)[0] * 100
    # find top predicted classes
    _, indices = torch.sort(output_data, descending=True)
    i = 0
    # print the top classes predicted by the model
    while confidences[indices[0][i]] > 0.5:
        class_idx = indices[0][i]
        print(
            "class:",
            classes[class_idx],
            ", confidence:",
            confidences[class_idx].item(),
            "%, index:",
            class_idx.item(),
        )
        i += 1
 
postprocess(output)

是時(shí)候測(cè)試我們的腳本了！我們的輸入圖像：

結(jié)果：

class: cup, confidence: 92.430747%, index: 968
class: espresso, confidence: 6.138075%, index: 967
class: coffee mug, confidence: 0.728557%, index: 504

2.將PyTorch模型轉(zhuǎn)換為ONNX格式

要轉(zhuǎn)換生成的模型，您只需要一行代碼torch.onnx.export，它需要以下參數(shù)：「預(yù)訓(xùn)練模型本身、與輸入數(shù)據(jù)大小相同的張量、ONNX 文件的名稱(chēng)、輸入和輸出名稱(chēng)」。

ONNX_FILE_PATH = 'resnet50.onnx'
torch.onnx.export(model, input, ONNX_FILE_PATH, input_names=['input'],
                  output_names=['output'], export_params=True)

要檢查模型轉(zhuǎn)換是否正常，請(qǐng)調(diào)用onnx.checker.check_model：

onnx_model = onnx.load(ONNX_FILE_PATH)
onnx.checker.check_model(onnx_model)

3. 可視化ONNX模型

現(xiàn)在，讓我們使用Netron可視化我們的 ONNX 圖。要啟動(dòng)它，請(qǐng)安裝：

python3 -m pip install netron

在命令行輸入netron并在瀏覽器中打開(kāi)http://localhost:8080/。您將看到完整的網(wǎng)絡(luò)圖。檢查輸入和輸出是否具有預(yù)期的大小。

4. 在TensorRT中初始化模型

現(xiàn)在是解析 ONNX 模型并初始化 TensorRT 「Context」和「Engine」的時(shí)候了。為此，我們需要?jiǎng)?chuàng)建一個(gè)Builder實(shí)例。Builder可以創(chuàng)建network并從該網(wǎng)絡(luò)生成engine（將針對(duì)您的平臺(tái)\硬件進(jìn)行優(yōu)化）。當(dāng)我們創(chuàng)建network時(shí)，我們可以通過(guò)標(biāo)志定義網(wǎng)絡(luò)的結(jié)構(gòu)，但在我們的例子中，使用默認(rèn)標(biāo)志就足夠了，這意味著所有張量都將具有隱式批次維度。通過(guò)network定義，我們可以創(chuàng)建一個(gè)Parser實(shí)例，最后解析我們的 ONNX 文件。

import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import tensorrt as trt
 
# logger to capture errors, warnings, and other information during the build and inference phases
TRT_LOGGER = trt.Logger()
 
def build_engine(onnx_file_path):
    # initialize TensorRT engine and parse ONNX model
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    parser = trt.OnnxParser(network, TRT_LOGGER)
     
    # parse ONNX
    with open(onnx_file_path, 'rb') as model:
        print('Beginning ONNX file parsing')
        parser.parse(model.read())
    print('Completed parsing of ONNX file')

可以配置一些engine參數(shù)，例如 TensorRT engine允許的最大內(nèi)存或設(shè)置 FP16 模式。我們還應(yīng)該指定批次的大小。

    # allow TensorRT to use up to 1GB of GPU memory for tactic selection
    builder.max_workspace_size = 1 << 30
    # we have only one image in batch
    builder.max_batch_size = 1
    # use FP16 mode if possible
    if builder.platform_has_fast_fp16:
        builder.fp16_mode = True

之后，我們可以生成「Engine」并創(chuàng)建可執(zhí)行文件「Context」。engine獲取輸入數(shù)據(jù)、執(zhí)行推理并發(fā)出推理輸出。

    # generate TensorRT engine optimized for the target platform
    print('Building an engine...')
    engine = builder.build_cuda_engine(network)
    context = engine.create_execution_context()
    print("Completed creating Engine")
 
    return engine, context

提示：初始化可能會(huì)花費(fèi)很多時(shí)間，因?yàn)?TensorRT 會(huì)嘗試找出在您的平臺(tái)上執(zhí)行網(wǎng)絡(luò)的最佳和更快的方式。要只執(zhí)行一次然后使用已經(jīng)創(chuàng)建的引擎，您可以序列化您的引擎。「序列化」引擎不能跨不同的 GPU 模型、平臺(tái)或 TensorRT 版本移植。引擎特定于它們所基于的確切硬件和軟件?？梢栽诖颂幷业礁嘈畔ⅲ篽ttps://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#serial_model_c。

5. 主函數(shù)

那么在 TensorRT 中進(jìn)行推理的完整流程會(huì)是什么樣子呢？讓我們看一下「主函數(shù)」。首先，讓我們解析模型并初始化engine和context：

def main():
    # initialize TensorRT engine and parse ONNX model
    engine, context = build_engine(ONNX_FILE_PATH)

當(dāng)我們擁有初始化引擎時(shí)，我們可以找出程序中輸入和輸出的維度。要知道我們可以分配輸入數(shù)據(jù)和輸出數(shù)據(jù)所需的內(nèi)存。在常見(jiàn)情況下，一個(gè)模型可以有一堆輸入和輸出，但在我們的例子中，我們知道我們只有一個(gè)輸入和一個(gè)輸出。

    # get sizes of input and output and allocate memory required for input data and for output data
    for binding in engine:
        if engine.binding_is_input(binding):  # we expect only one input
            input_shape = engine.get_binding_shape(binding)
            input_size = trt.volume(input_shape) * engine.max_batch_size * np.dtype(np.float32).itemsize     # in bytes
            device_input = cuda.mem_alloc(input_size)
        else:  # and one output
            output_shape = engine.get_binding_shape(binding)
            # create page-locked memory buffers (i.e. won't be swapped to disk)
            host_output = cuda.pagelocked_empty(trt.volume(output_shape) * engine.max_batch_size, dtype=np.float32)
            device_output = cuda.mem_alloc(host_output.nbytes)

CUDA 函數(shù)可以在流中異步調(diào)用。一個(gè)流中的所有命令將按順序執(zhí)行，但不同的流可以同時(shí)或亂序執(zhí)行它們的命令。當(dāng)您在未指定流的情況下執(zhí)行異步 CUDA 命令時(shí)，運(yùn)行時(shí)將使用默認(rèn)的空流。在我們的簡(jiǎn)單腳本中，我們將只創(chuàng)建一個(gè)流就足夠了。例如，在更復(fù)雜的情況下，您可以使用不同的流同時(shí)處理不同的圖像。

    # Create a stream in which to copy inputs/outputs and run inference.
    stream = cuda.Stream()

為了在 TensorRT 中獲得與在 PyTorch 中相同的結(jié)果，我們將為推理準(zhǔn)備數(shù)據(jù)并重復(fù)我們之前采取的所有預(yù)處理步驟。TensorRT 的 Python API 的主要好處是可以從 PyTorch 部分重用數(shù)據(jù)預(yù)處理和后處理。我們應(yīng)該做的唯一額外的事情是連續(xù)放置數(shù)據(jù)并盡可能使用page-locked memory。然后我們可以將該數(shù)據(jù)復(fù)制到 GPU 并將其用于推理。

    # preprocess input data
    host_input = np.array(preprocess_image("turkish_coffee.jpg").numpy(), dtype=np.float32, order='C')
    cuda.memcpy_htod_async(device_input, host_input, stream)

進(jìn)行推理并將結(jié)果從設(shè)備復(fù)制到主機(jī)：

    # run inference
    context.execute_async(bindings=[int(device_input), int(device_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(host_output, device_output, stream)
    stream.synchronize()

結(jié)果將存儲(chǔ)為host_output的一維數(shù)組。因此，在使用 PyTorch 部分的后處理來(lái)獲取人類(lèi)可讀的值之前，我們應(yīng)該對(duì)其進(jìn)行reshape。

    # postprocess results
    output_data = torch.Tensor(host_output).reshape(engine.max_batch_size, output_shape[0])
    postprocess(output_data)

就這樣！現(xiàn)在您可以啟動(dòng)腳本并對(duì)其進(jìn)行測(cè)試。

6. 精度測(cè)試

我們做了一些臨時(shí)測(cè)試，總結(jié)在下表中。

正如我們所見(jiàn)，預(yù)測(cè)的類(lèi)別匹配。置信度和 FP32 模式下幾乎相同（誤差小于 1e-05）。在 FP16 模式下錯(cuò)誤更大（~0.003），但它仍然足以獲得正確的預(yù)測(cè)。

請(qǐng)記住，不能保證您在使用不同的硬件、軟件甚至輸入圖片進(jìn)行測(cè)試時(shí)會(huì)遇到相同的精度。該精度可能取決于初始基準(zhǔn)決策，并且可能因不同的卡而不同。我們通過(guò)以下配置獲得這些結(jié)果：

Ubuntu 18.04.4, AMD? Ryzen 7 2700x eight-core processor × 16, GeForce RTX 2070 SUPER, TensorRT 6.0.1.5, CUDA 10.0

7. 使用 TensorRT 加速

為了比較 PyTorch 和 TensorRT 中的時(shí)間，我們不會(huì)測(cè)量模型的初始化時(shí)間，因?yàn)槲覀冎怀跏蓟艘淮?。所以我們將比較推理時(shí)間。在首次啟動(dòng)時(shí)，CUDA 會(huì)初始化并緩存一些數(shù)據(jù)，因此任何 CUDA 函數(shù)的首次調(diào)用都比平時(shí)慢。為了解決這個(gè)問(wèn)題，我們運(yùn)行推理幾次并獲得平均時(shí)間。我們擁有：

在我們的示例中，我們?cè)?FP16 模式下實(shí)現(xiàn)了 4-6 倍的加速，在 FP32 模式下實(shí)現(xiàn)了 2-3 倍的加速。

5、推薦閱讀

DETR即插即用 | RefineBox進(jìn)一步細(xì)化DETR家族的檢測(cè)框，無(wú)痛漲點(diǎn)

中科大提出PE-YOLO | 讓YOLO家族算法直擊黑夜目標(biāo)檢測(cè)

超越GIoU/DIoU/CIoU/EIoU | MPDIoU讓YOLOv7/YOLACT雙雙漲點(diǎn)，速度不減！