TensorRT part2 python version

最新推荐文章于 2025-06-28 09:11:37 发布

原创最新推荐文章于 2025-06-28 09:11:37 发布 · 1.2k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#tensorRT

本文是TensorRT系列的第二部分，主要介绍如何使用Python版本进行操作。内容包括回顾上文的TensorRT创建和推理流程图，然后详细阐述如何通过Python从TensorFlow模型转换到UFF格式，并利用官方例子进行TensorRT的Python应用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

TensorRT part2 python version

总结上文

在进入第二部分前，对第一部分的业务流程做一个总结：

创建流程图

这里写图片描述

推理流程图

这里写图片描述

pyversion

1.1 通过python使用TensorRT

只简单说明从tensorflow导入模型

import tensorrt as trt
from tensorrt.parsers import uffparser
import uff
#实现logging接口
G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)
#创建引擎
engine = uff.from_tensorflow_frozen_model(frozen_file, ["fc2/Relu"])
#创建解析器，和C++不一样，不是从builder中创建的解析器
parser = uffparser.create_uff_parser()
#注册输入和输出
parser.register_input("Placeholder", (1, 28, 28), 0)
parser.register_output("fc2/Relu")
# 将模型导入引擎
trt.utils.uff_to_trt_engine(G_LOGGER,uff_model,parser,MAX_BATCHSIZE,MAX_WORKSPACE)

1.2 将Frozen Graph转化为UFF

//通过命令行将.pb的Frozen Graph转化为UFF文件
convert-to-uff tensorflow -o name_of_output_uff_file --input-file name_of_input_pb_file -O name_of_output_tensor
//通过一下命令显示tensorflow的各层
convert-to-uff tensorflow --input-file name_of_input_pb_file -l

1.3 从tensorflow生成uff


#尝试导入uff模块
try:
    import uff
except ImportError:
        raise ImportError("""Please install the UFF Toolkit""")

OUTPUT_NAMES = ["fc2/Relu"]
UFF_OUTPUT_FILENAME = path + "/trained_lenet5.uff"

def train(xxx):
    #省略训练过程
    #。。。
    #。。。
    #训练完毕后

    #导出当前计算图的GraphDef部分
    graphdef = tf.get_default_graph().as_graph_def()
    #保存指定的节点，并将节点值保存为常数
    frozen_graph = tf.graph_util.convert_variables_to_constants(sess, graphdef, OUTPUT_NAMES)
    #删除不需要进行推理（inference）的节点
    return tf.graph_util.remove_training_nodes(frozen_graph)

#获得冻结计算图
frozen_graph = run_training(MNIST_DATASETS)
#将计算图转化为UFF形式
uff.from_tensorflow(graphdef=frozen_graph,
                        output_filename=UFF_OUTPUT_FILENAME,
                        output_nodes=OUTPUT_NAMES,
                        text=True)

1.4 通过python使用UFF（官方例子tf_to_trt）

    #截取代码

    #这个就是上面的train包装了一下
    #或者可以从文件读取，应该也能读取冻结计算图文件再进行转换
    tf_model = lenet5.learn()

    #导入UFF模型
    uff_model = uff.from_tensorflow(tf_model, ["fc2/Relu"])
    #Convert Tensorflow model to TensorRT model
    #创建一个解析器
    parser = uffparser.create_uff_parser()
    #注册输入和输出
    parser.register_input("Placeholder", (1, 28, 28), 0)
    parser.register_output("fc2/Relu")
    #记得实现logging接口
    G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.INFO)
    #构建引擎
    engine = trt.utils.uff_to_trt_engine(G_LOGGER,
                                              uff_model,
                                              parser,
                                              MAX_BATCHSIZE,
                                              MAX_WORKSPACE)
    #判定引擎不为空
    assert(engine)
    #获得执行上下文
    context = engine.create_execution_context()

    #进行推断
    print("\n| TEST CASE | PREDICTION |")
    for i in range(ITERATIONS):
        img, label = lenet5.get_testcase()
        img = img[0]
        label = label[0]
        #应该就只推断了一个
        out = infer(context, img, 1)
        print("|-----------|------------|")
        print("|     " + str(label) + "     |      " + str(np.argmax(out)) + "     |")

# API CHANGE: Try to generalize into a utils function
#Run inference on device
#推断过程
def infer(context, input_img, batch_size):
    #load engine
    #通过上下文载入引擎
    engine = context.get_engine()
    #不是很明白，这个判断的是什么
    assert(engine.get_nb_bindings() == 2)
    #create output array to receive data
    #创建输出的数组去接受数据
    dims = engine.get_binding_dimensions(1).to_DimsCHW()
    #计算数组的大小
    elt_count = dims.C() * dims.H() * dims.W() * batch_size
    #convert input data to Float32
    #将输入转化成float32
    input_img = input_img.astype(np.float32)
    #Allocate pagelocked memory
    #cuda申请页锁内存，即CPU和GPU共同拥有这块内存，省去显示传输过程
    output = cuda.pagelocked_empty(elt_count, dtype=np.float32)

    #alocate device memory
    #为输入输出申请GPU内存
    d_input = cuda.mem_alloc(batch_size * input_img.size * input_img.dtype.itemsize)
    d_output = cuda.mem_alloc(batch_size * output.size * output.dtype.itemsize)
    #返回的是地址？可能需要研究一下
    bindings = [int(d_input), int(d_output)]
    #创建CUDA流
    stream = cuda.Stream()

    #transfer input data to device
    #将数据拷贝到GPU
    cuda.memcpy_htod_async(d_input, input_img, stream)
    #execute model
    #推断流程
    context.enqueue(batch_size, bindings, stream.handle, None)
    #transfer predictions back
    #计算完数据同步回CPU
    cuda.memcpy_dtoh_async(output, d_output, stream)

    #return predictions
    return output