12.5 TensorRT API 使用
前言
在TensorRT的第二小節中，介紹了TensorRT模型計算圖的構建通常有三種方式，分別是TensorRT API、parser、訓練框架中trt工具。
三種方式的便捷程度與靈活度是依次遞增和遞減的，TensorRT API最靈活，可自由設定網路每一層，任意編輯模型的任何層，本節就介紹
TensorRT API的使用。
TensorRT API 簡介
TensorRT API是構建Network的一種方法，整個TensorRT的使用流程還是與第二節中介紹的一樣，整體流程如下圖所示
這裡不再贅述workflow內容，直接看TensorRT是如何構建network的，整體可分三個步驟：
- 創建network：network = builder.create_network()
- 逐層搭建網路：搭建網路層，並且賦予權重值
- 創建engine：engine = builder.build_engine(network, config)
由此可見，核心部分在於network提供了一系列網路層的搭建，這類似於pytorch中的nn.Module可以搭建各種網路層。
TensorRT官方提供的網路層可參考文檔
下面介紹採用TensorRT API搭建常見網路層的方法。
TensorRT API 網路層創建
卷積層
add_convolution(self: tensorrt.tensorrt.INetworkDefinition, input: tensorrt.tensorrt.ITensor, num_output_maps: int, kernel_shape: tensorrt.tensorrt.DimsHW, kernel: tensorrt.tensorrt.Weights, bias: tensorrt.tensorrt.Weights = None)→ tensorrt.tensorrt.IConvolutionLayer
Parameters：
- input – The input tensor to the convolution.
- num_output_maps – The number of output feature maps for the convolution.
- kernel_shape – The dimensions of the convolution kernel.
- kernel – The kernel weights for the convolution.
- bias – The optional bias weights for the convolution.
Returns：The new convolution layer, or None if it could not be created.
具體的創建案例代碼解釋如下：
- 設置輸入input：這個input如果是第一層的話，需要自行構建data；如果是中間層，則需要設置為上一層的輸出ITensor，例如pool2.get_output(0)返回的是一個ITensor物件。
- 設置卷積權重：在創建網路時，直接賦予卷積層的權重，這個權重通常通過wst檔獲取，是一個字典，key是自行標識的名稱，value是np.ndarray。
data = network.add_input('input', trt.float32, (3, 224, 224))
conv1 = network.add_convolution(input=data,
num_output_maps=64,
kernel_shape=(7, 7),
kernel=weight_map["conv1.weight"],
bias=trt.Weights())
fc1 = network.add_fully_connected(input=pool2.get_output(0),
num_outputs=OUTPUT_SIZE,
kernel=weight_map['fc.weight'],
bias=weight_map['fc.bias'])
Copy
BN層
在TensorRT API中沒有BN的實現，BN層需要通過add_scale層來實現，即通過減法、乘法的操作等價於實現BN層。
下面代碼是實現過程，核心是手動計算出需要shift和scale的數值，然後創建add_scale層。
def addBatchNorm2d(network, weight_map, input, layer_name, eps):
gamma = weight_map[layer_name + ".weight"]
beta = weight_map[layer_name + ".bias"]
mean = weight_map[layer_name + ".running_mean"]
var = weight_map[layer_name + ".running_var"]
var = np.sqrt(var + eps)
scale = gamma / var
shift = -mean / var * gamma + beta
return network.add_scale(input=input,
mode=trt.ScaleMode.CHANNEL,
shift=shift,
scale=scale)
Copy
啟動函數層
啟動函數層通過add_activation，其中提供了type屬性來設置不同的啟動函數。
relu3 = network.add_activation(ew1.get_output(0), type=trt.ActivationType.RELU)
Copy
type中可設置的類型有十多種，具體參考IActivationLayer的文檔
RELU : Rectified Linear activation
SIGMOID : Sigmoid activation
TANH : Hyperbolic Tangent activation
LEAKY_RELU : Leaky Relu activation: f(x) = x if x >= 0, f(x) = alpha * x if x < 0
ELU : Elu activation: f(x) = x if x >= 0, f(x) = alpha * (exp(x) - 1) if x < 0
SELU : Selu activation: f(x) = beta * x if x > 0, f(x) = beta * (alpha * exp(x) - alpha) if x <= 0
SOFTSIGN : Softsign activation: f(x) = x / (1 + abs(x))
SOFTPLUS : Softplus activation: f(x) = alpha * log(exp(beta * x) + 1)
CLIP : Clip activation: f(x) = max(alpha, min(beta, x))
HARD_SIGMOID : Hard sigmoid activation: f(x) = max(0, min(1, alpha * x + beta))
SCALED_TANH : Scaled Tanh activation: f(x) = alpha * tanh(beta * x)
THRESHOLDED_RELU : Thresholded Relu activation: f(x) = x if x > alpha, f(x) = 0 if x <= alpha
Copy
池化層
pool1 = network.add_pooling(input=relu1.get_output(0),
window_size=trt.DimsHW(3, 3),
type=trt.PoolingType.MAX)
Copy
全連接層
全連接層在v8.6.1之後進行了刪除，不再支援直接創建，需要用add_matrix_multiply、add_elementwise間接實現。
這個問題在官方文檔沒有指引，是通過多種方法找到的替代方案，這裡簡單記錄debug過程。
- 運行報錯：AttributeError: 'tensorrt.tensorrt.INetworkDefinition' object has no attribute 'add_fully_connected'
- 各種文檔找'add_fully_connected'，無果；
- 官方文檔查閱 ‘INetworkDefinition' 支持的方法，的確沒了add_fully_connected，猜測肯定有對應方法替換，梳理所有支持的方法，找到最相關的方法是：add_matrix_multiply
- TRT github倉庫搜索add_matrix_multiply， sample.py，發現了FC層實現方法:
具體方法如下：
def add_matmul_as_fc(net, input, outputs, w, b):
assert len(input.shape) >= 3
m = 1 if len(input.shape) == 3 else input.shape[0]
k = int(np.prod(input.shape) / m) # 輸入大小： 2048
assert np.prod(input.shape) == m * k
n = int(w.size / k) # 輸出大小： 1000
assert w.size == n * k
assert b.size == n
input_reshape = net.add_shuffle(input)
input_reshape.reshape_dims = trt.Dims2(m, k)
filter_const = net.add_constant(trt.Dims2(n, k), w)
mm = net.add_matrix_multiply(
input_reshape.get_output(0),
trt.MatrixOperation.NONE,
filter_const.get_output(0),
trt.MatrixOperation.TRANSPOSE,
)
bias_const = net.add_constant(trt.Dims2(1, n), b)
bias_add = net.add_elementwise(mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM)
output_reshape = net.add_shuffle(bias_add.get_output(0))
output_reshape.reshape_dims = trt.Dims4(m, n, 1, 1)
return output_reshape
Copy
V 8.6.1版本可用的方法如下：
fc1 = network.add_fully_connected(input=pool2.get_output(0),
num_outputs=OUTPUT_SIZE,
kernel=weight_map['fc.weight'],
bias=weight_map['fc.bias'])
Copy
以上是常見的網路層介紹，更多網路層參見文檔
tensorrtx
學習TensorRT API搭建網路模型，十分推薦根據tensorrtx的代碼資料學習。
tensorrtx是一個基於TensorRT API構建常用的網路模型庫，它完全基於TensorRT API搭建模型，這使得模型的構建具備高可調節性，也對理解、學習網路模型內部細節提供了很好的資料。
tensorrtx為模型提供了python和c++的代碼，使用tensorrtx的工作流程如下：
- 獲取訓練好的模型權重檔，例如pytorch的pt檔，tensorflow的ckpt檔，MXNet的.params+.json等
- 將權重文件匯出為.wts檔，wts檔TensorRT定義的一種模型參數檔案格式
- 創建network，逐網路層搭建模型，並基於wts檔賦予網路層權重值
- 搭建engine，進行推理或序列化保存到磁片。
wts 文件
這裡遇到一個新的檔案格式wts，wts是Weights的縮寫，是由Nvidia定義的一種模型參數檔案格式。
wts檔只包含模型的參數資料,不包含網路結構等其他資訊。網路結構需要另外定義，在這裡就需要用TensorRT API逐層定義。
通過使用wts檔分離參數和網路結構，可以方便地進行模型壓縮、量化等優化,也可以跨框架部署模型，同時還可保護模型的結構不被洩露。
ResNet50 推理
接下來演示如何採用TensorRT API搭建ResNet50，並完成推理，代碼來自tensorrx和pytorchx。
整體分兩步：
第一步，通過pytorchx的代碼獲得resnet.wts。在本教程，resnet.wts的生成可通過配套代碼實現
第二步，通過tensorrtx的代碼搭建TensorRT的network，並且創建engine進行推理。在本教程，推理代碼可通過配套代碼實現
wts檔生成
上述代碼已集成到配套章節代碼中，先看生成wts的核心代碼：
- wts檔第一行表明整個檔有多少個權重；
- 之後的每一行是一個權重名稱及權重值的二進位形式。
- f = open(path_wts, 'w')
- f.write("{}\n".format(len(net.state_dict().keys())))
- for k, v in net.state_dict().items():
- print('key: ', k)
- print('value: ', v.shape)
- vr = v.reshape(-1).cpu().numpy()
- f.write("{} {}".format(k, len(vr)))
- for vv in vr:
- f.write(" ")
- f.write(struct.pack(">f", float(vv)).hex())
- f.write("\n")
Copy
以下是一個wts案例：
10
conv1.weight 150 be40ee1b bd20bab8 bdc4bc53 .......
conv1.bias 6 bd327058 .......
conv2.weight 2400 3c6f2220 3c693090 ......
Copy
network創建
接下來參照第一份配套代碼，將engine的構建從直接載入檔改為基於network創建的形式
回顧之前直接讀取磁片上的engine檔，可以直接得到engine。
with open(model_path, 'rb') as ff, trt.Runtime(logger) as runtime:
engine = runtime.deserialize_cuda_engine(ff.read())
Copy
接下來看network逐層創建的過程，主workflow是一樣的，需要創建各模組，這裡有兩個重點。
- trt推薦採用顯式batch，因此create_network的時候需要設置1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
- resnet50的創建，在函數build_model_by_trt_api中實現
def init_model(model_path):
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
network = build_model_by_trt_api(network, model_path)
config = builder.create_builder_config()
engine = builder.build_engine(network, config)
context = engine.create_execution_context()
context.set_input_shape("data", [1, 3, 224, 224]) # 綁定輸入張量的形狀
return context, engine
Copy
build_model_by_trt_api函數中就是上文提到的TensorRT API的各網路層創建介面，這裡就不一一贅述，到這裡，整體的流程梳理完成。
def build_model_by_trt_api(network, model_path):
weight_map = load_wts_file(model_path) # 權重字典
data = network.add_input("data", trt.float32, (1, 3, 224, 224))
assert data
conv1 = network.add_convolution(input=data,
num_output_maps=64,
kernel_shape=(7, 7),
kernel=weight_map["conv1.weight"],
bias=trt.Weights())
Copy
最終能得到如下圖片，表明基於TensorRT API創建網路模型並推理是成功的。
tensorrtx還提供了非常多常見的網路模型，如果有需要使用 TensorRT API創建網路模型，建議優選tensorrtx中尋找參考案例。
小結
本小節介紹TensorRT API進行trt模型創建的流程及案例。首先介紹了TensorRT API創建網路層的方法及常用介面，隨後介紹學習TensorRT API 非常好的代碼庫——tensorrtx，最後以一個resnet50的推理案例介紹基於TensorRT API 搭建模型的全流程。
通過本小節，可以瞭解TensorRT API 創建模型的過程與概念，更多案例及用法，請閱讀：
- python api文檔：https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/index.htm
- tensorrtx：https://github.com/wang-xinyu/tensorrtx
留言列表
