TensorFlow 笔记02-mnist 的 tensorRT 实现，从 .npz 文件中加载参数进行推理

● 代码，tf 卷积神经网络，将训练好的参数保存为 .npz 文件给 tensorRT 用

# tf 模型搭建和训练部分同上一篇博客
tfArg = {}
for i in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):      # 遍历全局元素集合，全部放进 tfArg 中来
    tfArg[i.name] = sess.run(i)
tfArg['testX']=mnist[2].images                                  # 补上 mnist 的测试数据
tfArg['testY']=mnist[2].labels

np.savez(pbFilePath + 'tfArg.npz',**tfArg)                      # 保存 tfArg 为 .npz 文件
sess.close()

● 代码，将前面一模一样的神经网络用 trt 重写一遍，加载训练好的参数来推理

import numpy as np
import tensorflow as tf
import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda
import input_data
from datetime import datetime as dt

pbFilePath = "tempFile/"

# 网络基础设施
iGpu = 0
print("GPU in use:", cuda.Device(iGpu).name())
cuda.Device(iGpu).make_context()
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network()
builder.max_batch_size = 64
builder.max_workspace_size = 0 << 20

# 读取参数
para = np.load(pbFilePath + 'tfArg.npz')
w1= para['w1:0'].transpose((3, 2, 0, 1)).reshape(-1)                            # NHWC -> NCHW，所有权重都要 reshape(-1) 压成 1 维
b1 = para['b1:0']
w2 = para['w2:0'].transpose((3, 2, 0, 1)).reshape(-1)
b2 = para['b2:0']
w3 = para['w3:0'].reshape(7,7,64,1024).transpose((3, 2, 0, 1)).reshape(-1)
b3 = para['b3:0']
w4 = para['w4:0'].reshape(1024,10).transpose((1,0)).reshape(-1)
b4 = para['b4:0']
testX = para['testX']                                                           # 测试数据
testY = para['testY']

# 建立网络
batchSize = 64
data = network.add_input("data", trt.DataType.FLOAT, (batchSize, 1, 28, 28))    # 输入层，batchSize 张 1 通道 28 行 28 列

h1 = network.add_convolution(data, 32, (5, 5), w1, b1)                          # 卷积 1，指定输出特征数，窗口高宽，权重值（隐式转换为 trt.Weigfhts）
h1.stride = (1, 1)                                                              # 外侧补充指定跨步和光环
h1.padding = (2, 2)
h1Act = network.add_activation(h1.get_output(0), trt.ActivationType.RELU)       # 激活层，指定激活类型

h1Pool = network.add_pooling(h1Act.get_output(0), trt.PoolingType.MAX, (2, 2))  # 池化层，指定池化类型，窗口高宽
h1Pool.stride = (2, 2)
h1Pool.padding = (0, 0)
               
h2 = network.add_convolution(h1Pool.get_output(0), 64, (5, 5), w2, b2)          # 卷积 2
h2.stride = (1, 1)
h2.padding = (2, 2)
h2Act = network.add_activation(h2.get_output(0), trt.ActivationType.RELU)

h2Pool = network.add_pooling(h1Act.get_output(0), trt.PoolingType.MAX, (2, 2))  # 池化 2
h2Pool.stride = (2, 2)
h2Pool.padding = (0, 0)


h3 = network.add_fully_connected(h2Pool.get_output(0), 1024, w3, b3)            # 全连接层，指定输出特征数，权重值
h3Act = network.add_activation(h3.get_output(0), trt.ActivationType.RELU)

h4 = network.add_fully_connected(h3Act.get_output(0), 10, w4, b4)               # 全连接层 2
y = network.add_softmax(h4.get_output(0))                                       # softmax 层

network.mark_output(y.get_output(0))                                            # 指定输出层
engine = builder.build_cuda_engine(network)                                     # 建立 engine

# 申请内存
h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=np.float32)  
h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=np.float32)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)

# 流和上下文
stream = cuda.Stream()
context = engine.create_execution_context()

# 测试
print( "%s, start!" %( dt.now()) )
acc = 0
nTest = len(para['testX'])
for i in range(nTest // batchSize):                                             # 向下取整，尾巴可能没测完       
    h_input = para['testX'][i*batchSize:(i+1)*batchSize].reshape(-1,1,28,28)
    
    cuda.memcpy_htod_async(d_input, h_input, stream)                            # 数据拷贝
    
    context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)  # 执行内核
    
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    
    stream.synchronize()                                                        # 同步，否则 yy 是 全零矩阵
    
    yy = np.argmax(h_output.reshape(engine.get_binding_shape(1)),1).reshape(-1)
    label = np.argmax(para['testY'][i*batchSize:(i+1)*batchSize],1)    
    acc += np.sum( ( yy == label ).astype(np.int) )
    
cuda.Context.pop()                                                              # 删除上下文
print( "%s, acc = %f" %( dt.now(), acc/(len(para['testX'])) ) )

▶ 总结 tensorRT 的关键步骤（包含 engine 的读写，避免每次都新建 engine 浪费时间）

import tensorflow as tf
import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda

iGpu = 0
cuda.Device(iGpu).make_context()                # 设备上下文
logger = trt.Logger(trt.Logger.WARNING)         # 创建 logger

trtFilePath = "./densenetEngine.trt"            # 读取现成的 engine 序列文件，否则现场生成一个 engine 并序列化保存为文件
if os.path.isfile(trtFilePath) and not DEBUG:
    with open(trtFilePath, 'rb') as f:
        engineStr = f.read()
else:
    builder = trt.Builder(logger)               # 创建 builder
    builder.max_batch_size     = 64
    builder.max_workspace_size = 200 << 20
    builder.fp16_mode          = True           # 是否使用 float16
    
    network = builder.create_network()          # 创建 network

    h0 = network.add_input("h0", ...)           # 开始建网
    
    ...

    y = network.add_...
    
    network.mark_output(y.get_output(0))        # 标记输出节点
    
    engine = builder.build_cuda_engine(network) # 建立 engine，最容易失败的位置

    if engine == None:
        print("build engine failed!")
        return None        
    
    engineStr = engine.serialize()              # 创建序列化的 engine 并写入文件中，方便下次直接取用
    with open(trtFilePath, 'wb') as f:
        f.write(engineStr)

runtime = trt.Runtime(logger)                                               # 利用运行时环境读取序列化的 engine（现场创建 engine 的可以跳过这步）
engine  = runtime.deserialize_cuda_engine(engineStr)
context = engine.create_execution_context()                                 # 创建内核上下文（区别于设备上下文）
stream  = cuda.Stream()                                                     # 创建流（可选）

hIn = cuda.pagelocked_empty(engine.get_binding_shape(0), dtype=np.float32)  # 使用无初始化的页锁定内存，指定尺寸（隐式转换为 trt.volume）和数据类型，也可用 np.empty 等来申请一般内存
hOut = cuda.pagelocked_empty(engine.get_binding_shape(1), dtype=np.float32) # engine.get_binding_shape 的 (0) 和 (1) 分别等于network 的输入和输出节点尺寸
dIn = cuda.mem_alloc(h_input.nbytes)                                        # 申请设备内存，使用主机内存的大小
dOut = cuda.mem_alloc(h_output.nbytes)

cuda.memcpy_htod_async(d_input, h_input, stream)                            # 异步数据拷贝
#cuda.memcpy_htod(d_input, data)                                            # 非异步数据拷贝
context.execute_async(batchSize, bindings=[int(dIn), int(dOut)], stream_handle=stream.handle)  # 异步执行内核
context.execute(batchSize, bindings=[int(dIn), int(dOut)])                  # 非异步执行内核
cuda.memcpy_dtoh_async(hOut, dOut, stream)

stream.synchronize()                                                        # 同步

context = None                                                              # 清空内核上下文和 engine
engine  = None
cuda.Context.pop()                                                          # 关闭设备上下文

▶ 留坑，使用 convert_to_uff.py 将保存的 .pb 模型转化为 .uff 模型，方便 tendorRT 直接加载和使用，不用再在 tenorRT 中重建。中间遇到一些问题，尚未成功。