3. Working With TensorRT Using The Python API

以下部分将重点介绍使用 Python API 可以执行 TensorRT 用户的目标和任务。这些部分主要讨论在没有任何框架的情况下使用 Python API。示例部分提供了进一步的详细信息,并在适当情况下链接到下面。

假设你从一个训练过的模型开始。本章将介绍以下使用TensorRT的必要步骤:

  • 从模型中创建 TensorRT 网络定义
  • 调用 TensorRT builder 从网络创建优化的运行时引擎
  • 序列化和反序列化引擎,以便在运行时快速重新创建引擎
  • 向引擎提供数据以执行推理

Python API vs C++ API

从本质上说,C++ API 和 Python API 在支持您的需求方面应该接近相同。Python API 的主要优点是数据预处理和后处理很容易使用,因为您可以使用各种库,如 NumPy 和 SciPy。

C++ API 应该用于安全非常重要的场合,例如在汽车中。有关 C++ API 的更多信息,请参见使用 C++ API 处理 TensorRT

有关如何使用 Python 优化性能的更多信息,请参见如何优化我的 Python 性能?来自最佳实践指南。

3.1. Importing TensorRT Into Python

  1. 导入 TensorRT:
import tensorrt as trt
  1. 实现一个日志接口,通过该接口 TensorRT 报告错误、警告和信息消息。下面的代码展示了如何实现日志记录接口。在这种情况下,我们抑制了信息消息,只报告警告和错误。TensorRT Python 绑定中包含了一个简单的日志记录器。
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

3.2. Creating A Network Definition In Python

使用 TensorRT 执行推理的第一步是从您的模型创建一个 TensorRT 网络。最简单的方法是使用 TensorRT 解析器库导入模型(请参阅使用 Python 中的解析器导入模型、使用 Python 从 Caffe 导入模型、使用 Python 从 TensorFlow 导入模型和使用 Python 从 ONNX 导入模型),该库支持以下格式的序列化模型:

  • Caffe (both BVLC and NVCaffe)
  • ONNX 1.0 and 1.1, and
  • UFF (used for TensorFlow)

另一种选择是直接使用 TensorRT 网络 API 定义模型(请参阅使用Python API从头创建网络定义)。这需要您进行少量的 API 调用来定义网络图中的每一层,并为模型的训练参数实现您自己的导入机制。

注意: TensorRT Python API 仅适用于 x86_64 平台。更多信息请参见深度学习 SDK 文档- TensorRT 工作流。

3.2.1. Creating A Network Definition From Scratch Using The Python API

在创建网络时,必须首先定义引擎并创建用于推理的构建器对象。Python API 用于从网络 API 创建网络和引擎。网络定义引用用于向网络添加各种层。有关使用 Python API 创建网络和引擎的更多信息,请参见 network_api_pytorch_mnist 示例。

下面的代码演示了如何创建一个具有输入、卷积、池、完全连接、激活和 SoftMax 层的简单网络。

# Create the builder and network
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network:
    # Configure the network layers based on the weights provided. In this case, the weights are imported from a pytorch model. 
    # Add an input layer. The name is a string, dtype is a TensorRT dtype, and the shape can be provided as either a list or tuple.
    input_tensor = network.add_input(name=INPUT_NAME, dtype=trt.float32, shape=INPUT_SHAPE)

    # Add a convolution layer
    conv1_w = weights['conv1.weight'].numpy()
    conv1_b = weights['conv1.bias'].numpy()
    conv1 = network.add_convolution(input=input_tensor, num_output_maps=20, kernel_shape=(5, 5), kernel=conv1_w, bias=conv1_b)
    conv1.stride = (1, 1)

    pool1 = network.add_pooling(input=conv1.get_output(0), type=trt.PoolingType.MAX, window_size=(2, 2))
    pool1.stride = (2, 2)
    conv2_w = weights['conv2.weight'].numpy()
    conv2_b = weights['conv2.bias'].numpy()
    conv2 = network.add_convolution(pool1.get_output(0), 50, (5, 5), conv2_w, conv2_b)
    conv2.stride = (1, 1)

    pool2 = network.add_pooling(conv2.get_output(0), trt.PoolingType.MAX, (2, 2))
    pool2.stride = (2, 2)

    fc1_w = weights['fc1.weight'].numpy()
    fc1_b = weights['fc1.bias'].numpy()
    fc1 = network.add_fully_connected(input=pool2.get_output(0), num_outputs=500, kernel=fc1_w, bias=fc1_b)

    relu1 = network.add_activation(fc1.get_output(0), trt.ActivationType.RELU)

    fc2_w = weights['fc2.weight'].numpy()
    fc2_b = weights['fc2.bias'].numpy()
    fc2 = network.add_fully_connected(relu1.get_output(0), OUTPUT_SIZE, fc2_w, fc2_b)

    fc2.get_output(0).name =OUTPUT_NAME
    network.mark_output(fc2.get_output(0))

3.2.2. Importing A Model Using A Parser In Python

要使用解析器导入模型,需要执行以下高级步骤:

  • 创建TensorRTbuilder和网络。
  • 为特定格式创建TensorRT解析器。
  • 使用解析器解析导入的模型并填充网络。

有关这些步骤和示例代码的示例,请参见使用 Python 从 Caffe 导入、使用 Python 从 TensorFlow 导入和使用 Python 从 ONNX 导入。
构建器必须在网络之前创建,因为它是网络的工厂。不同的解析器有不同的网络输出标记机制。有关更多信息,请参见 UFF 解析器 APICaffe 解析器 APIONNX 解析器 API

3.2.3. Importing From Caffe Using Python

下面的步骤说明了如何使用 Caffe 解析器和 Python API 直接导入 Caffe 模型。有关更多信息,请参阅 introductory_parser_samples 示例。

  1. import TensorRT:
import tensorrt as trt
  1. 定义数据类型。在本例中,我们将使用float32。
datatype = trt.float32
  1. 另外,定义一些路径。更改以下路径,以反映您在示例中所包含的模型中所放置的位置:
deploy_file = 'data/mnist/mnist.prototxt'
model_file = 'data/mnist/mnist.caffemodel'
  1. 创建 builder, network, 和 parser:
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.CaffeParser() as parser:
model_tensors = parser.parse(deploy=deploy_file, model=model_file, network=network, dtype=datatype)

解析器返回 model_tensors,它是一个表,包含从张量名称到 ITensor 对象的映射。

3.2.4. Importing From TensorFlow Using Python

下面的步骤说明了如何使用 UffParser 和 Python API 直接导入 TensorFlow 模型。这个示例可以在 <site-packages>/tensorrt/samples/python/end_to_end_tensorflow_mnist 目录中找到。有关更多信息,请参见 end_to_end_tensorflow_mnist Python 示例。

  1. Import TensorRT:
import tensorrt as trt
  1. 为TensorFlow模型创建一个冻结的TensorFlow模型。 关于将TensorFlow模型冻结到流中的说明可以在 Freezing A TensorFlow Graph 中找到。

  2. 使用 UFF 转换器将冻结的 tensorflow 模型转换为 UFF 文件。通常,这很简单:

convert-to-uff frozen_inference_graph.pb

根据您如何安装 TensorRT,转换到 uff 实用程序可能不会安装在您的系统路径中。在这种情况下,直接调用底层 Python 脚本。它应该位于 UFF 模块的 bin 目录中;例如, ~/.local/lib/python2.7/site-packages/uff/bin/convert_to_uff.py

要找到 UFF 模块的位置,运行如下命令 python -c "import uff; print(uff.__path__)"

或者,您可以使用 UFF Paser API 并直接转换 TensorFlow GraphDef。

  1. 定义一些路径。更改以下路径,以反映您将示例中包含的模型放置在何处:
model_file = '/data/mnist/minist.uff'
  1. 创建 builder, network, 和 parser:
with builder = trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
        parser.register_input("Placeholder", (1, 28, 28))
        parser.register_output("fc2/Relu")
        parser.parse(model_file, network)

3.2.5. Importing From ONNX Using Python

限制:由于 ONNX 格式发展很快,您可能会遇到模型版本和解析器版本之间的版本不匹配。TensorRT 5.0.0 附带的 ONNX 解析器支持 ONNX IR (Intermediate Representation)版本 0.0.3,opset 版本 7。

一般来说,ONNX 解析器的新版本是向后兼容的,因此,遇到由早期版本的 ONNX 导出器生成的模型文件不会造成问题。当更改不向后兼容时,可能会有一些例外。在这种情况下,将早期的 ONNX 模型文件转换为稍后支持的版本。有关此主题的更多信息,请参见 ONNX Model Opset Version Converter

用户模型也有可能是由一个导出工具生成的,该工具支持比 TensorRT 附带的 ONNX 解析器支持的更晚的 opset。在这种情况下,检查发布到 GitHub onnx-tensorrt 的最新版本是否支持所需的版本。有关更多信息,请参见 yolov3_onnx

支持的版本由 onnx_trt_backend.cpp 中的 BACKEND_OPSET_VERSION 变量定义。从 GitHub 下载并构建 ONNX TensorRT 解析器的最新版本。构建的说明可以在这里找到: TensorRT backend for ONNX.

下面的步骤说明如何使用 OnnxParser 和 Python API 直接导入 ONNX 模型。有关更多信息,请参见 introductory_parser_samples Python示例。

  1. import TensorRT:
import tensorrt as trt
  1. 创建 build, network, and parser:
with builder = trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
with open(model_path, 'rb') as model:
parser.parse(model.read())

3.2.6. Importing From PyTorch And Other Frameworks

在 PyTorch 中使用 TensorRT (或任何其他具有 NumPy 兼容权重的框架)涉及到使用 TensorRT API 复制网络体系结构(请参阅 Creating A Network Definition From Scratch Using The Python API),然后从 PyTorch 复制权重。有关更多信息,请参见 Working With PyTorch And Other Frameworks

注意:在 ubuntu14.04 和 CentOS 上,同时加载 torch 模块和 TensorRT 可能会导致 segmentation faults。

要执行推理,请遵循 Performing Inference In Python 的概述说明。

3.3. Building An Engine In Python

构建器的一个功能是搜索其 CUDA 内核目录以获得最快的可用实现,因此必须使用相同的 GPU 来构建优化引擎将运行的 GPU。

构建器具有许多属性,您可以设置这些属性以控制网络应运行的精度,以及自动调整参数,例如TensorRT在确定哪个最快时(多次迭代会导致更长的运行时间,但是对噪声的敏感性较小)应该为每个内核计时多少次。您还可以查询构建器以找出硬件本机支持的混合精度类型。

两个特别重要的属性是最大批处理大小(maximum batch size )和最大工作区大小(maximum workspace size)。

最大批大小指定 TensorRT 要优化的批大小。在运行时,可以选择较小的批处理大小。

层算法通常需要临时工作区。此参数限制网络中任何层可以使用的最大大小。如果提供的划痕不够,TensorRT 可能无法找到给定层的实现。

有关在Python中构建引擎的更多信息,请参见 introductory_parser_samples 示例。

  1. 使用构建对象构建 engine:
builder.max_batch_size = max_batch_size
builder.max_workspace_size = 1 <<  20 # This determines the amount of memory available to the builder when building an optimized engine and should generally be set as high as possible.
with trt.Builder(TRT_LOGGER) as builder:
with builder.build_cuda_engine(network) as engine:
# Do inference here.
  1. 执行推理。要执行推理,请遵循在 Performing Inference In Python 中概述的说明。

3.4. Serializing A Model In Python

当您序列化时,您将引擎转换为一种格式,以便稍后存储并用于推理。要用于推理,只需对引擎进行反序列化即可。序列化和反序列化是可选的。由于从网络定义创建引擎可能会耗费大量时间,因此您可以避免在每次应用程序重新运行时重新构建引擎,方法是对引擎进行一次序列化,并在进行推断时对其进行反序列化。因此,在构建引擎之后,用户通常希望对其进行序列化,以便以后使用。

从这里开始,您可以序列化引擎,也可以直接使用引擎进行推理。在将模型用于推理之前,序列化和反序列化是一个可选步骤——如果需要,引擎对象可以直接用于推理。

注意:序列化引擎不能跨平台或 TensorRT 版本移植。引擎特定于它们所构建的 GPU 模型(除了平台和 TensorRT 版本之外)。

  1. 将模型序列化为 modelstream:
serialized_engine = engine.serialize()
  1. 反序列化 modelstream 以执行推理。反序列化需要创建一个运行时对象:
with trt.Runtime(TRT_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(serialized_engine)

最后一个参数是使用定制层应用程序的插件层工厂,它是可选的。可以在 Extending TensorRT With Custom Layers 中找到更多细节。

也可以将序列化引擎保存到文件中,并从文件中读回:

  1. 序列化引擎并写入文件:
with open(“sample.engine”, “wb”) as f:
        f.write(engine.serialize())
  1. 从文件中读取引擎并反序列化:
with open(“sample.engine”, “rb”) as f, trt.Runtime(TRT_LOGGER) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())

3.5. Performing Inference In Python

下面的步骤说明了如何在Python中执行推断,现在您已经有了一个引擎。

  1. 为输入和输出分配一些主机和设备缓冲区:
# Determine dimensions and create page-locked memory buffers (i.e. won't be swapped to disk) to hold host inputs/outputs.
        h_input = cuda.pagelocked_empty(engine.get_binding_shape(0).volume(), dtype=np.float32)
        h_output = cuda.pagelocked_empty(engine.get_binding_shape(1).volume(), dtype=np.float32)
        # Allocate device memory for inputs and outputs.
        d_input = cuda.mem_alloc(h_input.nbytes)
        d_output = cuda.mem_alloc(h_output.nbytes)
        # Create a stream in which to copy inputs/outputs and run inference.
        stream = cuda.Stream()
  1. 创建一些空间来存储中间激活值。由于引擎包含网络定义和经过训练的参数,因此需要额外的空间。这些是在执行上下文中持有的:
with engine.create_execution_context() as context:
        # Transfer input data to the GPU.
        cuda.memcpy_htod_async(d_input, h_input, stream)
        # Run inference.
        context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
        # Transfer predictions back from the GPU.
        cuda.memcpy_dtoh_async(h_output, d_output, stream)
        # Synchronize the stream
        stream.synchronize()
        # Return the host output. 
return h_output

引擎可以有多个执行上下文,允许一组权重用于多个重叠的推理任务。例如,您可以在并行 CUDA 流中处理图像,每个流使用一个引擎和一个上下文。每个上下文将在与引擎相同的 GPU 上创建。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 200,527评论 5 470
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 84,314评论 2 377
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 147,535评论 0 332
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,006评论 1 272
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 62,961评论 5 360
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,220评论 1 277
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,664评论 3 392
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,351评论 0 254
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,481评论 1 294
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,397评论 2 317
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,443评论 1 329
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,123评论 3 315
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,713评论 3 303
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,801评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,010评论 1 255
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 42,494评论 2 346
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,075评论 2 341

推荐阅读更多精彩内容