TensorRT 加速 Tensorflow 实现 Mnist 数据集分类详解

项目介绍

在 Jetson Nano 上使用 TensorRT 为 Mnist 数据集的推理过程进行加速。使用的 .py 文件可以在下面这个路径中找到。

cd /usr/src/tensorrt/python/end_to_end_tensorflow_mnist

版本介绍

在 Jetson Nano 上安装完 Jetpack 之后，发现它自带了 TensorRT，这里给出我使用的版本：

• Jetpack：4.3；
• CUDA：10.0；
• cuDNN：7.6.3；
• Python：3.6；
• TensorRT：7.0；
• Tensorflow：1.15.0。

model.py 代码说明

import tensorflow as tf
import numpy as np

# 数据集处理
def process_dataset():
    (x_train, y_train),(x_test, y_test) = tf.keras.datasets.mnist.load_data()  # 下载数据集
    x_train, x_test = x_train / 255.0, x_test / 255.0  # 数据归一化

    NUM_TRAIN = 60000  # 训练集中样本个数
    NUM_TEST = 10000  # 测试集中样本个数
    x_train = np.reshape(x_train, (NUM_TRAIN, 28, 28, 1))  # 将训练集样本的 shape reshape 成 [28, 28, 1]
    x_test = np.reshape(x_test, (NUM_TEST, 28, 28, 1))  # 将测试集样本的 shape reshape 成 [28, 28, 1]
    return x_train, y_train, x_test, y_test

# 构建 LeNet5 模型
def create_model():
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.InputLayer(input_shape=[28, 28, 1]))
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(512, activation=tf.nn.relu))
    model.add(tf.keras.layers.Dense(10, activation=tf.nn.softmax))
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])  # 编译模型--优化方法用 Adam，损失函数使用 sparse_categorical_crossentropy
    # 【注】独热编码用 categorical_crossentropy，数字编码用 sparse_categorical_crossentropy
    return model

# 保存模型到 .pb 文件
def save(model, filename):
	sess = tf.keras.backend.get_session()  # 返回 TF 会话
	input_graph_def = sess.graph.as_graph_def()  # 获取 GraphDef 对象
    output_names = model.output.op.name  # 得到输出名称
    frozen_graph = tf.graph_util.convert_variables_to_constants(sess, input_graph_def, [output_names])  # 将 GraphDef 对象中的变量转化为常量
    frozen_graph = tf.graph_util.remove_training_nodes(frozen_graph)  # 删除不需要进行推理的节点
    # 保存序列化后的 GraphDef 到磁盘
    with open(filename, "wb") as ofile:
        ofile.write(frozen_graph.SerializeToString())

def main():
    x_train, y_train, x_test, y_test = process_dataset()  # 加载样本
    model = create_model()  # 创建模型
    model.fit(x_train, y_train, epochs = 5, verbose = 1)  # 训练模型
    model.evaluate(x_test, y_test)  # 测试模型
    save(model, filename="models/lenet5.pb")  # 将模型冻结为 .pb 文件，这里的 models 文件夹需要被提前建好

if __name__ == '__main__':
    main()

以上代码中除了保存模型部分之外都是一些常规操作，这里不再赘述。对于 save(model, filename) 函数，它的具体步骤如下：

• 获取 GraphDef 对象;
• 找到需要导出的节点 output_names ;
• 使用 convert_variables_to_constants 方法将 GraphDef 对象中的变量转化为常量，并返回一个新的 GraphDef；
• 删除不需要进行推理的节点；
• 序列化保存到磁盘。

将 .pb 文件转换为 .uff 文件

有了有了上一步得到的 .pb 模型之后，我们可以利用 TensorRT 提供的转换工具将该模型转换为 .uff 格式模型：

python3 /usr/lib/python3.6/dist-packages/uff/bin/convert_to_uff.py --input_file models/lenet5.pb

其中的 python3.6 需要对应自己的 python 版本号，models/lenet5.pb 是 .pb 文件相对于当前位置的路径。

sample.py 代码说明

from random import randint
from PIL import Image
import numpy as np

import pycuda.driver as cuda
import pycuda.autoinit

import tensorrt as trt

import sys, os
sys.path.insert(1, os.path.join(sys.path[0], ".."))
import common  # 这个文件也在 /usr/src/tensorrt/python/end_to_end_tensorflow_mnist 中

# 实现日志记录接口，TensorRT 通过该接口报告错误、警告和信息性消息
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)  # 抑制信息性消息，仅报告警告和错误

# 设置网络信息和样本信息
class ModelData(object):
    MODEL_FILE = "lenet5.uff"  # 网络模型名称
    INPUT_NAME ="input_1"  # 网络输入层名称
    INPUT_SHAPE = (1, 28, 28)  # 样本 shape，这里的顺序是（channels, height, width）
    OUTPUT_NAME = "dense_1/Softmax"  # 网络输出层名称

# 加载并处理测试图片
def load_normalized_test_case(data_paths, pagelocked_buffer, case_num=randint(0, 9)):  # 因为样本名称是从 0.pgm 到 9.pgm，所以 case_num=randint(0, 9)
    [test_case_path] = common.locate_files(data_paths, [str(case_num) + ".pgm"])
    # Flatten 该图像成为一个 1 维数组，然后归一化，并 copy 到主机的 pagelocked 内存中
    img = np.array(Image.open(test_case_path)).ravel()
    np.copyto(pagelocked_buffer, 1.0 - img / 255.0)
    return case_num

def main():
    data_paths, _ = common.find_sample_data(description="Runs an MNIST network using a UFF model file", subfolder="mnist")  # 这里返回的是存放样本的文件夹的路径
    model_path = os.environ.get("MODEL_PATH") or os.path.join(os.path.dirname(__file__), "models")  # 这里返回的是存放模型 .uff 文件的文件夹的路径
    model_file = os.path.join(model_path, ModelData.MODEL_FILE)  # 这里返回的是模型 .uff 文件的路径

    # 建立引擎
    parser = trt.UffParser()  # 创建 parser
    builder = trt.Builder(TRT_LOGGER)  # 创建 builder
    builder.max_workspace_size = common.GiB(1)  # 为 builder 指定一个最大的 workspace 空间
    network = builder.create_network()  # 创建 network
    
    parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)  # 注册网络的输入
    parser.register_output(ModelData.OUTPUT_NAME)  # 注册网络的输出
    parser.parse(model_file, network)  # 解析网络
    
    engine = builder.build_cuda_engine(network)  # 构建引擎
    
    inputs, outputs, bindings, stream = common.allocate_buffers(engine)  # 将数据从主机传输到 GPU
    
    with engine.create_execution_context() as context:  # 通过引擎得到 context
        case_num = load_normalized_test_case(data_paths, pagelocked_buffer=inputs[0].host)  # 读取测试样本，并归一化
        [output] = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)  # 对测试样本进行推理，返回一个列表，此例中列表内只有一个元素
        # 如果有多个 input/output 节点，则使用 common.do_inference_v2 函数
        pred = np.argmax(output)
        print("Test Case: " + str(case_num))
        print("Prediction: " + str(pred))

if __name__ == '__main__':
    main()

【注】因为创建引擎是很耗时的，所以强烈建议将引擎保存下来：

serialized_engine = engine.serialize()  # 序列化引擎
# 写入文件
with open(engine_path, 'wb') as f:
	f.write(serialized_engine)

这样一来，我们就不用重复创建引擎了，只需要加载即可：

# 从文件读取引擎并反序列化
with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
	return runtime.deserialize_cuda_engine(f.read())

在推断过程中，我们使用了 common.do_inference 函数，它的具体步骤如下：

• 将输入的数据从主机传输到 GPU；
• 通过异步执行去做推断；
• 将推断的结果从 GPU 传回主机；
• 将 stream 同步。

其代码如下：

def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
    # 将数据移动到 GPU
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # 执行异步推断
    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
    # 将结果从 GPU 写回到 host 端（GPU --> 主机）
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # 同步 stream
    stream.synchronize()
    # 返回主机的输出结果
    return [out.host for out in outputs]

在终端中运行 sample.py 文件，最终得到推断结果：

报错

如果运行的时候报如下错的话：

RuntimeError: __iter__() is only supported inside of tf.function or when eager execution is enabled.

则在导入 Tensorflow 库之后加一句：

tf.enable_eager_execution()

来开启 eager 模式。