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 中重建。中間遇到一些問題，尚未成功。