從 TensorRT 7 . 0 開始， Universal Framework Format（ UFF ）被棄用。在本文中，您將學習如何使用新的 TensorFlow -ONNX- TensorRT 工作流部署經過 TensorFlow 培訓的深度學習模型。圖 1 顯示了 TensorRT 的高級工作流。
首先，使用任何框架訓練網絡。網絡訓練后，批量大小和精度是固定的（精度為 FP32 、 FP16 或 INT8 ）。訓練好的模型被傳遞給 TensorRT 優化器，優化器輸出一個優化的運行時（也稱為計劃）。. plan 文件是 TensorRT 引擎的序列化文件格式。計劃文件需要反序列化才能使用 TensorRT 運行時運行推斷。
要優化在 TensorFlow 中實現的模型，只需將模型轉換為 ONNX 格式，并使用 TensorRT 中的 ONNX 解析器解析模型并構建 TensorRT 引擎。圖 2 顯示了高級 ONNX 工作流。
在本文中，我們將討論如何使用 ONNX 工作流創建一個 TensorRT 引擎，以及如何從 TensorRT 引擎運行推理。更具體地說，我們演示了從 Keras 或 TensorFlow 中的模型到 ONNX 的端到端推理，以及使用 ResNet-50 、語義分段和 U-Net 網絡的 TensorRT 引擎。最后，我們將解釋如何在其他網絡上使用此工作流。
下載本文中的代碼示例。 下載 TensorFlow -onnx- TensorRT 后 – 代碼 tar . gz 文件，您還應該從 Cityscapes dataset scripts repo 下載 labels.py ，并將其與其他腳本放在同一個文件夾中。
ONNX 概述
ONNX 是機器學習和深度學習模型的開放格式。它允許您將不同框架（如 TensorFlow 、 PyTorch 、 MATLAB 、 Caffe 和 Keras ）的深度學習和機器學習模型轉換為單一格式。
它定義了一組通用的運算符、深入學習的通用構建塊集和通用文件格式。它提供計算圖的定義以及內置運算符。可能有一個或多個輸入或輸出的 ONNX 節點列表形成一個無環圖。
ResNet ONNX 工作流示例
在這個例子中，我們展示了如何在兩個不同的網絡上使用 ONNX 工作流并創建一個 TensorRT 引擎。第一個網絡是 ResNet-50 。
工作流包括以下步驟：
- 將 TensorFlow / Keras 模型轉換為. pb 文件。
- 將. pb 文件轉換為 ONNX 格式。
- 創建 TensorRT 引擎。
- 從 TensorRT 引擎運行推斷。
將模型轉換為. pb
第一步是將模型轉換為. pb 文件。以下代碼示例將 ResNet-50 模型轉換為. pb 文件：
import tensorflow as tf
import keras
from tensorflow.keras.models import Model
import keras.backend as K
K.set_learning_phase(0)
def keras_to_pb(model, output_filename, output_node_names):
"""
This is the function to convert the Keras model to pb.
Args:
model: The Keras model.
output_filename: The output .pb file name.
output_node_names: The output nodes of the network. If None, then
the function gets the last layer name as the output node.
"""
# Get the names of the input and output nodes.
in_name = model.layers[0].get_output_at(0).name.split(':')[0]
if output_node_names is None:
output_node_names = [model.layers[-1].get_output_at(0).name.split(':')[0]]
sess = keras.backend.get_session()
# The TensorFlow freeze_graph expects a comma-separated string of output node names.
output_node_names_tf = ','.join(output_node_names)
frozen_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
sess.graph_def,
output_node_names)
sess.close()
wkdir = ''
tf.train.write_graph(frozen_graph_def, wkdir, output_filename, as_text=False)
return in_name, output_node_names
# load the ResNet-50 model pretrained on imagenet
model = keras.applications.resnet.ResNet50(include_top=True, weights='imagenet', input_tensor=None, input_shape=None, pooling=None, classes=1000)
# Convert the Keras ResNet-50 model to a .pb file
in_tensor_name, out_tensor_names = keras_to_pb(model, "models/resnet50.pb", None) 除了 Keras ，您還可以從以下位置下載 ResNet-50 ：
- 深度學習示例 GitHub 存儲庫：提供最新的深度學習示例網絡。您還可以看到 ResNet-50 分支，它包含一個腳本和方法來訓練 ResNet-50v1 . 5 模型。
- NVIDIA NGC 型號 ：它有預訓練模型的檢查點列表。例如，在 ResNet-50v1 . 5 上搜索 TensorFlow ，并從 Download 頁面獲取最新的檢查點。
將. pb 文件轉換為 ONNX
第二步是將. pb 模型轉換為 ONNX 格式。為此，首先安裝 tf2onnx 。
安裝 tf2onnx 后，有兩種方法可以將模型從. pb 文件轉換為 ONNX 格式。第二種方法是使用命令行。運行以下命令：
python -m tf2onnx.convert --input /Path/to/resnet50.pb --inputs input_1:0 --outputs probs/Softmax:0 --output resnet50.onnx
從 ONNX 創建 TensorRT 引擎
要從 ONNX 文件創建 TensorRT 引擎，請運行以下命令：
import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.WARNING) trt_runtime = trt.Runtime(TRT_LOGGER) def build_engine(onnx_path, shape = [1,224,224,3]): """ This is the function to create the TensorRT engine Args: onnx_path : Path to onnx_file. shape : Shape of the input of the ONNX file. """ with trt.Builder(TRT_LOGGER) as builder, builder.create_network(1) as network, trt.OnnxParser(network, TRT_LOGGER) as parser: builder.max_workspace_size = (256 << 20) with open(onnx_path, 'rb') as model: parser.parse(model.read()) network.get_input(0).shape = shape engine = builder.build_cuda_engine(network) return engine def save_engine(engine, file_name): buf = engine.serialize() with open(file_name, 'wb') as f: f.write(buf) def load_engine(trt_runtime, plan_path): with open(engine_path, 'rb') as f: engine_data = f.read() engine = trt_runtime.deserialize_cuda_engine(engine_data) return engine
此代碼應保存在引擎. py 文件，稍后在文章中使用。
此代碼示例包含以下變量：
- 最大工作區大小： 在執行時
ICudaEngine可以使用的最大 GPU 臨時內存。
構建器創建一個空網絡（ builder.create_network() ）， ONNX 解析器將 ONNX 文件解析到網絡（ parser.parse(model.read()) ）。您可以為網絡（ network.get_input(0).shape = shape ）設置輸入形狀，然后生成器將創建引擎（ engine = builder.build_cuda_engine(network) ）。要創建引擎，請運行以下代碼示例：
import engine as eng import argparse from onnx import ModelProto import tensorrt as trt engine_name = “resnet50.plan” onnx_path = "/path/to/onnx/result/file/" batch_size = 1 model = ModelProto() with open(onnx_path, "rb") as f: model.ParseFromString(f.read()) d0 = model.graph.input[0].type.tensor_type.shape.dim[1].dim_value d1 = model.graph.input[0].type.tensor_type.shape.dim[2].dim_value d2 = model.graph.input[0].type.tensor_type.shape.dim[3].dim_value shape = [batch_size , d0, d1 ,d2] engine = eng.build_engine(onnx_path, shape= shape) eng.save_engine(engine, engine_name)
在這個代碼示例中，首先從 ONNX 模型獲取輸入形狀。接下來，創建引擎，然后將引擎保存在. plan 文件中。
運行來自 TensorRT 引擎的推理：
TensorRT 引擎在以下工作流中運行推理：
- 為 GPU 中的輸入和輸出分配緩沖區。
- 將數據從主機復制到 GPU 中分配的輸入緩沖區。
- 在 GPU 中運行推理。
- 將結果從 GPU 復制到主機。
- 根據需要重塑結果。
下面的代碼示例詳細解釋了這些步驟。此代碼應保存在推理. py 文件，稍后將在本文中使用。
import tensorrt as trt import pycuda.driver as cuda import numpy as np import pycuda.autoinit def allocate_buffers(engine, batch_size, data_type): """ This is the function to allocate buffers for input and output in the device Args: engine : The path to the TensorRT engine. batch_size : The batch size for execution time. data_type: The type of the data for input and output, for example trt.float32. Output: h_input_1: Input in the host. d_input_1: Input in the device. h_output_1: Output in the host. d_output_1: Output in the device. stream: CUDA stream. """ # Determine dimensions and create page-locked memory buffers (which won't be swapped to disk) to hold host inputs/outputs. h_input_1 = cuda.pagelocked_empty(batch_size * trt.volume(engine.get_binding_shape(0)), dtype=trt.nptype(data_type)) h_output = cuda.pagelocked_empty(batch_size * trt.volume(engine.get_binding_shape(1)), dtype=trt.nptype(data_type)) # Allocate device memory for inputs and outputs. d_input_1 = cuda.mem_alloc(h_input_1.nbytes) d_output = cuda.mem_alloc(h_output.nbytes) # Create a stream in which to copy inputs/outputs and run inference. stream = cuda.Stream() return h_input_1, d_input_1, h_output, d_output, stream def load_images_to_buffer(pics, pagelocked_buffer): preprocessed = np.asarray(pics).ravel() np.copyto(pagelocked_buffer, preprocessed) def do_inference(engine, pics_1, h_input_1, d_input_1, h_output, d_output, stream, batch_size, height, width): """ This is the function to run the inference Args: engine : Path to the TensorRT engine pics_1 : Input images to the model. h_input_1: Input in the host d_input_1: Input in the device h_output_1: Output in the host d_output_1: Output in the device stream: CUDA stream batch_size : Batch size for execution time height: Height of the output image width: Width of the output image Output: The list of output images """ load_images_to_buffer(pics_1, h_input_1) with engine.create_execution_context() as context: # Transfer input data to the GPU. cuda.memcpy_htod_async(d_input_1, h_input_1, stream) # Run inference. context.profiler = trt.Profiler() context.execute(batch_size=1, bindings=[int(d_input_1), int(d_output)]) # Transfer predictions back from the GPU. cuda.memcpy_dtoh_async(h_output, d_output, stream) # Synchronize the stream stream.synchronize() # Return the host output. out = h_output.reshape((batch_size,-1, height, width)) return out
為第一個輸入行和輸出行確定兩個維度。您可以在主機（ h_input_1 、 h_output ）中創建頁鎖定內存緩沖區。然后，為輸入和輸出分配與主機輸入和輸出相同大小的設備內存（ d_input_1 ， d_output ）。下一步是創建 CUDA 流，用于在設備和主機分配的內存之間復制數據。
在這個代碼示例中，在 do_inference 函數中，第一步是使用 load_images_to_buffer 函數將圖像加載到主機中的緩沖區。然后將輸入數據傳輸到 GPU （ cuda.memcpy_htod_async(d_input_1, h_input_1, stream) ），并使用 context.execute 運行推理。最后將結果從 GPU 復制到主機（ cuda.memcpy_dtoh_async(h_output, d_output, stream) ）。
ONNX 工作流語義分割實例
在本文 基于 TensorRT 3 的自主車輛快速 INT8 推理 中，作者介紹了一個語義分割模型的 UFF 工作流過程。
在本文中，您將使用類似的網絡來運行 ONNX 工作流來進行語義分段。該網絡由一個基于 VGG16 的編碼器和三個使用反褶積層實現的上采樣層組成。網絡在 城市景觀數據集 上經過大約 40000 次迭代訓練
有多種方法可以將 TensorFlow 模型轉換為 ONNX 文件。一種方法是 ResNet50 部分中解釋的方法。 Keras 也有自己的 Keras 到 ONNX 文件轉換器。有時， TensorFlow -to-ONNX 不支持某些層，但 Keras-to-ONNX 轉換器支持這些層。根據 Keras 框架和使用的層類型，您可能需要在轉換器之間進行選擇。
在下面的代碼示例中，使用 Keras-to-ONNX 轉換器將 Keras 模型直接轉換為 ONNX 。下載預先訓練的語義分段文件 semantic_segmentation.hdf5 。
import keras
import tensorflow as tf
from keras2onnx import convert_keras
def keras_to_onnx(model, output_filename):
onnx = convert_keras(model, output_filename)
with open(output_filename, "wb") as f:
f.write(onnx.SerializeToString())
semantic_model = keras.models.load_model('/path/to/semantic_segmentation.hdf5')
keras_to_onnx(semantic_model, 'semantic_segmentation.onnx') 圖 3 顯示了網絡的體系結構。
與前面的示例一樣，使用下面的代碼示例創建語義分段引擎。
import engine as eng from onnx import ModelProto import tensorrt as trt engine_name = 'semantic.plan' onnx_path = "semantic.onnx" batch_size = 1 model = ModelProto() with open(onnx_path, "rb") as f: model.ParseFromString(f.read()) d0 = model.graph.input[0].type.tensor_type.shape.dim[1].dim_value d1 = model.graph.input[0].type.tensor_type.shape.dim[2].dim_value d2 = model.graph.input[0].type.tensor_type.shape.dim[3].dim_value shape = [batch_size , d0, d1 ,d2] engine = eng.build_engine(onnx_path, shape= shape) eng.save_engine(engine, engine_name)
要測試模型的輸出，請使用 城市景觀數據集 。要使用城市景觀，必須具有以下功能： sub_mean_chw 和 color_map 。這些函數也用于 post ， 基于 TensorRT 3 的自主車輛快速 INT8 推理 .
在下面的代碼示例中， sub_mean_chw 用于從圖像中減去平均值作為預處理步驟， color_map 是從類 ID 到顏色的映射。后者用于可視化。
import numpy as np from PIL import Image import tensorrt as trt import labels # from cityscapes evaluation script import skimage.transform TRT_LOGGER = trt.Logger(trt.Logger.WARNING) trt_runtime = trt.Runtime(TRT_LOGGER) MEAN = (71.60167789, 82.09696889, 72.30508881) CLASSES = 20 HEIGHT = 512 WIDTH = 1024 def sub_mean_chw(data): data = data.transpose((1, 2, 0)) # CHW -> HWC data -= np.array(MEAN) # Broadcast subtract data = data.transpose((2, 0, 1)) # HWC -> CHW return data def rescale_image(image, output_shape, order=1): image = skimage.transform.resize(image, output_shape, order=order, preserve_range=True, mode='reflect') return image def color_map(output): output = output.reshape(CLASSES, HEIGHT, WIDTH) out_col = np.zeros(shape=(HEIGHT, WIDTH), dtype=(np.uint8, 3)) for x in range(WIDTH): for y in range(HEIGHT): if (np.argmax(output[:, y, x] )== 19): out_col[y,x] = (0, 0, 0) else: out_col[y, x] = labels.id2label[labels.trainId2label[np.argmax(output[:, y, x])].id].color return out_col
下面的代碼示例是上一個示例的其余代碼。必須先運行上一個塊，因為需要定義的函數。使用這個例子比較 Keras 模型和 TensorRT 引擎 semantic . plan 文件的輸出，然后可視化這兩個輸出。根據需要替換占位符 /path/to/semantic_segmentation.hdf5 和 input_file_path 。
import engine as eng
import inference as inf
import keras
import tensorrt as trt
input_file_path = ‘munster_000172_000019_leftImg8bit.png’
onnx_file = "semantic.onnx"
serialized_plan_fp32 = "semantic.plan"
HEIGHT = 512
WIDTH = 1024
image = np.asarray(Image.open(input_file_path))
img = rescale_image(image, (512, 1024),order=1)
im = np.array(img, dtype=np.float32, order='C')
im = im.transpose((2, 0, 1))
im = sub_mean_chw(im)
engine = eng.load_engine(trt_runtime, serialized_plan_fp32)
h_input, d_input, h_output, d_output, stream = inf.allocate_buffers(engine, 1, trt.float32)
out = inf.do_inference(engine, im, h_input, d_input, h_output, d_output, stream, 1, HEIGHT, WIDTH)
out = color_map(out)
colorImage_trt = Image.fromarray(out.astype(np.uint8))
colorImage_trt.save(“trt_output.png”)
semantic_model = keras.models.load_model('/path/to/semantic_segmentation.hdf5')
out_keras= semantic_model.predict(im.reshape(-1, 3, HEIGHT, WIDTH))
out_keras = color_map(out_keras)
colorImage_k = Image.fromarray(out_keras.astype(np.uint8))
colorImage_k.save(“keras_output.png”) 圖 4 顯示了實際圖像和實際情況，以及 Keras 的輸出與 TensorRT 引擎的輸出的對比。如您所見， TensorRT 發動機的輸出與 Keras 的類似。
在其他網絡上試試
現在您可以在其他網絡上嘗試 ONNX 工作流。有關分段網絡的好例子的更多信息，請參閱 GitHub 上的 具有預訓練主干的分割模型 。
作為一個例子，我們用一個 ONNX 網絡來說明如何使用。本例中的網絡是來自 segmentation_models 庫的 U-Net 。在這里，我們只加載模型，而沒有對其進行訓練。您可能需要在首選數據集上訓練這些模型。
關于這些網絡的一個重要點是，當您加載這些網絡時，它們的輸入層大小如下所示：（ None ， None ， None ， 3 ）。要創建一個 TensorRT 引擎，您需要一個輸入大小已知的 ONNX 文件。在將此模型轉換為 ONNX 之前，請通過為其輸入指定大小來更改網絡，然后將其轉換為 ONNX 格式。
例如，從這個庫（ segmentation _ models ）加載 U-Net 網絡并為其輸入指定大小（ 244 、 244 、 3 ）。在為推理創建了 TensorRT 引擎之后，做一個與語義分段類似的轉換。根據應用程序和數據集的不同，可能需要使用不同的顏色映射。
import segmentation_models as sm import keras from keras2onnx import convert_keras from engine import * onnx_path = 'unet.onnx' engine_name = 'unet.plan' batch_size = 1 CHANNEL = 3 HEIGHT = 224 WIDTH = 224 model = sm.Unet() model._layers[0].batch_input_shape = (None, 224,224,3) model = keras.models.clone_model(model) onx = convert_keras(model, onnx_path) with open(onnx_path, "wb") as f: f.write(onx.SerializeToString()) shape = [batch_size , HEIGHT, WIDTH, CHANNEL] engine = build_engine(onnx_path, shape= shape) save_engine(engine, engine_name)
我們之前提到的另一種下載方式是從 vz6 下載。它有一個預先訓練模型的檢查點列表。例如，您可以在 TensorFlow 中搜索 UNet ，然后轉到 Download 頁面以獲取最新的檢查點。
總結
在這篇文章中，我們解釋了如何使用 TensorFlow-to-ONNX-to-TensorRT 工作流來部署深度學習應用程序，并給出了幾個示例。第一個例子是 ResNet-50 上的 ONNX- TensorRT ，第二個例子是在 Cityscapes 數據集上訓練的基于 英偉達數據中心深度學習產品性能 的語義分割。在文章的最后，我們演示了如何在其他網絡上應用這個工作流。
