Python API 构建 TensorRT 加速模型的步骤详解
TensorRT（TensorRT是一种高性能神经网络推理（模型推断）引擎，主要用于在生产环境中部署深度学习模型。）是NVIDIA深度学习SDK中的一部分，是一种高效的深度学习推断加速库。TensorRT 可以将深度学习推理模型构建成一个高度优化的计算图形，用于部署到不同的 NVIDIA GPU 设备上，从而实现比原始模型更快、更小、更加节能的运行效果。
Python API 是基于TensorRT C++ API的一个封装，可以极大的简化TensorRT模型构建的流程。本文将详细介绍如何利用Python API构建TensorRT模型。
步骤
TensorRT模型构建的主要流程如下：
- 定义网络结构
- 加载训练好的模型
- 创建推理引擎
- 序列化引擎
- 反序列化引擎
定义网络结构
定义网络结构主要是用来在TensorRT中构建计算图。TensorRT采用了Caffe的模型概念来定义网络结构。在Python API中，可以使用caffe_pb2模块中定义的网络相关类来定义网络结构。
例如，下面的代码定义了一个简单的网络结构：
import tensorrt as trt
import numpy as np
import os
# Define network
G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)
builder = trt.Builder(G_LOGGER)
network = builder.create_network()
data = network.add_input(name="data", dtype=trt.float32, shape=(3, 224, 224))
fc1_w = np.random.rand(4096, 25088).astype(np.float32)
fc1_b = np.random.rand(4096).astype(np.float32)
fc1_layer = network.add_fully_connected(data, 4096, fc1_w, fc1_b)
relu1_layer = network.add_activation(fc1_layer.get_output(0), trt.infer.ActivationType.RELU)
fc2_w = np.random.rand(4096, 4096).astype(np.float32)
fc2_b = np.random.rand(4096).astype(np.float32)
fc2_layer = network.add_fully_connected(relu1_layer.get_output(0), 4096, fc2_w, fc2_b)
relu2_layer = network.add_activation(fc2_layer.get_output(0), trt.infer.ActivationType.RELU)
fc3_w = np.random.rand(1000, 4096).astype(np.float32)
fc3_b = np.random.rand(1000).astype(np.float32)
fc3_layer = network.add_fully_connected(relu2_layer.get_output(0), 1000, fc3_w, fc3_b)
softmax = network.add_softmax(fc3_layer.get_output(0))
softmax.get_output(0).set_name("prob")
network.mark_output(tensor=softmax.get_output(0))
加载训练好的模型
加载训练好的模型是指将训练好的模型导入到TensorRT中。目前，TensorRT支持Caffe和UFF（Universal Framework Format）格式的模型。下面的示例中，我们将使用Caffe格式的模型（.prototxt 文件和 .caffemodel文件）。
import tensorrt as trt
import numpy as np
import os
# Define network
G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)
builder = trt.Builder(G_LOGGER)
network = builder.create_network()
# Load trained model
model_path = "/path/to/model"
deploy_file = os.path.join(model_path, "deploy.prototxt")
model_file = os.path.join(model_path, "model.caffemodel")
parser = trt.CaffeParser()
parser.set_plugin_factory_ext(my_plugin_factory)
blob_name_to_tensor = parser.parse(deploy_file, model_file, network, trt.DataType.FLOAT)
# Set input and output
data = blob_name_to_tensor["data"]
network.mark_input(data.name)
在上述示例中，我们使用了CaffeParser类来解析Caffe格式的模型，并将模型导入到TensorRT网络中。parse()函数的第一个参数是.prototxt文件的路径，第二个参数是.caffemodel文件的路径。在将模型导入到网络中后，我们需要定义输入和输出张量。
创建推理引擎
创建推理引擎是TensorRT模型构建的核心步骤。推理引擎是一个高度优化的计算图形，用于快速推断神经网络模型。在Python API中，可以使用ICudaEngine类来创建推理引擎。
import tensorrt as trt
import numpy as np
import os
# Define network
G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)
builder = trt.Builder(G_LOGGER)
network = builder.create_network()
# Load trained model
model_path = "/path/to/model"
deploy_file = os.path.join(model_path, "deploy.prototxt")
model_file = os.path.join(model_path, "model.caffemodel")
parser = trt.CaffeParser()
parser.set_plugin_factory_ext(my_plugin_factory)
blob_name_to_tensor = parser.parse(deploy_file, model_file, network, trt.DataType.FLOAT)
# Set input and output
data = blob_name_to_tensor["data"]
network.mark_input(data.name)
# Create engine
builder.max_batch_size = 16
builder.max_workspace_size = 2 << 30
engine = builder.build_cuda_engine(network)
在上述示例中，我们使用Builder类来创建推理引擎。在构建推理引擎时，我们需要指定最大批处理大小和最大的工作空间大小。
序列化引擎
序列化是将推理引擎保存到磁盘上的过程。保存后，我们可以在不需要重新构建网络的情况下重复使用推理引擎。在Python API中，可以使用ICudaEngine类的serialize()函数来序列化推理引擎。
import tensorrt as trt
import numpy as np
import os
# Define network
G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)
builder = trt.Builder(G_LOGGER)
network = builder.create_network()
# Load trained model
model_path = "/path/to/model"
deploy_file = os.path.join(model_path, "deploy.prototxt")
model_file = os.path.join(model_path, "model.caffemodel")
parser = trt.CaffeParser()
parser.set_plugin_factory_ext(my_plugin_factory)
blob_name_to_tensor = parser.parse(deploy_file, model_file, network, trt.DataType.FLOAT)
# Set input and output
data = blob_name_to_tensor["data"]
network.mark_input(data.name)
# Create engine
builder.max_batch_size = 16
builder.max_workspace_size = 2 << 30
engine = builder.build_cuda_engine(network)
# Serialize engine
engine_path = "/path/to/engine"
serialized_engine = engine.serialize()
f = open(engine_path, "wb")
f.write(serialized_engine)
f.close()
在上述示例中，我们将序列化后的推理引擎保存到了/engine目录下。
反序列化引擎
反序列化是将已序列化的推理引擎重新构建成推理引擎对象的过程。在Python API中，可以使用ICudaEngine类的deserialize()函数来反序列化推理引擎。
import tensorrt as trt
# Deserialize engine
engine_path = "/path/to/engine"
f = open(engine_path, "rb")
serialized_engine = f.read()
f.close()
runtime = trt.infer.create_infer_runtime(G_LOGGER)
engine = runtime.deserialize_cuda_engine(serialized_engine)
在上述示例中，我们使用反序列化函数deserialize_cuda_engine()来反序列化推理引擎。
示例
以下是一个简单的示例，演示如何使用TensorRT Python API构建一个简单的卷积神经网络模型。这个模型包含一个卷积层、一个池化层和一个全连接层，用于对MNIST手写数字进行识别。
import tensorrt as trt
import numpy as np
G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)
def add_conv_layer(network, input, out_channels, kernel_size):
conv_w = np.ones((out_channels, kernel_size, kernel_size)).astype(np.float32)
conv_b = np.zeros((out_channels,)).astype(np.float32)
conv_layer = network.add_convolution(input, out_channels, (kernel_size, kernel_size), conv_w, conv_b)
return conv_layer
def add_pool_layer(network, input, pool_type, window_size):
pool_layer = network.add_pooling(input, pool_type, (window_size, window_size))
return pool_layer
def add_fc_layer(network, input, out_features):
fc_w = np.ones((input.get_output(0).shape[-1], out_features)).astype(np.float32)
fc_b = np.zeros((out_features,)).astype(np.float32)
fc_layer = network.add_fully_connected(input, out_features, fc_w, fc_b)
return fc_layer
def build_network():
builder = trt.Builder(G_LOGGER)
network = builder.create_network()
input_tensor = network.add_input(name="input", dtype=trt.float32, shape=(1, 28, 28))
conv_layer = add_conv_layer(network, input_tensor, out_channels=32, kernel_size=5)
pool_layer = add_pool_layer(network, conv_layer.get_output(0), trt.PoolingType.AVERAGE, 2)
relu_layer = network.add_activation(pool_layer.get_output(0), trt.infer.ActivationType.RELU)
fc_layer = add_fc_layer(network, relu_layer.get_output(0), out_features=10)
softmax_layer = network.add_softmax(fc_layer.get_output(0))
output_tensor = softmax_layer.get_output(0)
output_tensor.set_name("output")
network.mark_output(output_tensor)
return builder, network
def build_engine():
builder, network = build_network()
builder.max_batch_size = 1
builder.max_workspace_size = 2 << 30
engine = builder.build_cuda_engine(network)
return engine
def main():
engine = build_engine()
runtime = trt.infer.create_infer_runtime(G_LOGGER)
context = engine.create_execution_context()
# Generate random input
input_data = np.random.rand(1, 28, 28).astype(np.float32)
# Allocate device memory
input_shape = engine.get_binding_shape(0)
input_size = trt.volume(input_shape) * engine.max_batch_size * np.dtype(np.float32).itemsize
d_input = cuda.mem_alloc(input_size)
# Copy input data to device
cuda.memcpy_htod(d_input, input_data)
# Allocate output memory
output_shape = engine.get_binding_shape(1)
output_size = trt.volume(output_shape) * engine.max_batch_size * np.dtype(np.float32).itemsize
d_output = cuda.mem_alloc(output_size)
# Run inference
bindings = [int(d_input), int(d_output)]
context.execute_async_v2(bindings=bindings)
# Copy output data from device
output_data = np.empty_like(output_data)
cuda.memcpy_dtoh(output_data, d_output)
# Print output
print(output_data)
if __name__ == "__main__":
main()
在上述示例中，我们构建了一个简单的卷积神经网络模型，用于对MNIST手写数字进行识别。模型中包含一个卷积层、一个池化层和一个全连接层。我们使用Builder类来构建引擎，使用create_execution_context()函数来创建执行上下文，使用execute_async_v2()函数来运行推断，最后从设备上复制输出数据到主机并打印。
总结
在本文中，我们详细介绍了如何使用TensorRT Python API构建一个模型，包括定义网络结构、加载训练好的模型、创建推理引擎、序列化引擎和反序列化引擎等步骤。我们还提供了一个简单的示例，演示了如何使用TensorRT Python API构建一个简单的卷积神经网络模型。
本站文章如无特殊说明，均为本站原创，如若转载，请注明出处：Python api构建tensorrt加速模型的步骤详解 - Python技术站
微信扫一扫 
支付宝扫一扫 
