極智AI | TensorRT 中 Layer 和 Tensor 的區別

語言: CN / TW / HK

  一起養成寫作習慣!這是我參與「掘金日新計劃 · 6 月更文挑戰」的第9天,點擊查看活動詳情

歡迎關注我的公眾號 [極智視界],獲取我的更多筆記分享

  大家好,我是極智視界,本文介紹一下 TensorRT 中 Layer 和 Tensor 的區別。

  TensorRT 模型構建由一層層的 Layer 計算節點 和 一層層的 Tensor 數據節點組織而成,那麼什麼是計算節點,什麼是數據節點,以及計算節點和數據節點有什麼區別呢,這裏我們進行介紹。

1 區分 Layer 和 Tensor

  區分 Layer 和 Tensor,可以理解為區分 計算節點數據節點 ,這和 ONNX 中的 node 和 tensor 的概念比較類似,用過 ONNX 的同學可能比較清楚一些。來看 TensorRT 中如何定義一個 Layer,以及如何從一個 Layer 獲取它的輸出 Tensor:

```python

output is a layer

oneLayer = network.add_identity(inputTensor)

get tensor from the layer

oneTensor = oneLayer.get_output(0)

take the tensor into next layer

nextLayer = network.add_identity(oneTensor)
```

  可以看到 TensorRT 中,添加一個 Layer 或者説 構造一個 Layer,一般就用 network.add_xxx,而從 Layer 獲取一個 Tensor 一般可以用 xxxLayer.get_output(i) 。兩者的區別,用通俗一些的話,Layer 構建了一個骨架,而 Tensor 填滿了它。

2 Layer 和 Tensor 的常用方法

  這裏介紹 TensorRT 中 Layer 和 Tensor 的一些常用方法,以便於在構建網絡或者調試的時候有更多的方法可以去定位 / 解決問題。

Layer 的常用成員和方法

  • oneLayer.name = 'one': 獲取或指定 Layer 的名字;
  • oneLayer.type: 獲取該層的種類;
  • oneLayer.precision: 指定該層計算精度 (需配合 builder.strict_type_constraints);
  • oneLayer.get_output(i): 獲取該層第i個輸出張量;

Tensor 的常用成員和方法: - oneTensor.name = 'one': 獲取指定 tensor 的名字; - oneTensor.shape: 獲取 tensor 的形狀,可用於 print 檢查或作為後續層的參數; - oneTensor.dtype: 獲取或設定 tensor 的數據類型 (可用於配合 identity 層實現數據類型轉換);

3 示例演示

 從 Network 中打印所有層和張量的信息

  • 外層循環遍歷所有 Layer;
  • 內層循環遍歷該 Layer 的所有 input / output;

  來看示例:

```python import os import numpy as np import tensorrt as trt

logger = trt.Logger(trt.Logger.ERROR) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) profile = builder.create_optimization_profile() config = builder.create_builder_config() config.max_workspace_size = 3 << 30 print(config.flags)

inputTensor = network.add_input('inputT0', trt.DataType.FLOAT, [-1, 1, 28, 28]) profile.set_shape(inputTensor.name, (1, 1, 28, 28), (4, 1, 28, 28), (8, 1, 28, 28)) config.add_optimization_profile(profile)

w = np.random.rand(32,1,5,5).astype(np.float32).reshape(-1) b = np.random.rand(32).astype(np.float32).reshape(-1) _0 = network.add_convolution_nd(inputTensor, 32, [5, 5], w, b) _0.padding_nd = [2, 2] _1 = network.add_activation(_0.get_output(0), trt.ActivationType.RELU) _2 = network.add_pooling_nd(_1.get_output(0), trt.PoolingType.MAX, [2, 2]) _2.stride_nd = [2, 2]

w = np.random.rand(64,32,5,5).astype(np.float32).reshape(-1) b = np.random.rand(64).astype(np.float32).reshape(-1) _3 = network.add_convolution_nd(_2.get_output(0), 64, [5, 5], w, b) _3.padding_nd = [2, 2] _4 = network.add_activation(_3.get_output(0), trt.ActivationType.RELU) _5 = network.add_pooling_nd(_4.get_output(0), trt.PoolingType.MAX, [2, 5]) _5.stride_nd = [2, 2]

_6 = network.add_shuffle(_5.get_output(0)) _6.first_transpose = (0, 2, 3, 1) _6.reshape_dims = (-1, 64 * 7 * 7, 1, 1)

w = np.random.rand(1024,64 * 7 * 7).astype(np.float32).reshape(-1) b = np.random.rand(1024).astype(np.float32).reshape(-1) _7 = network.add_fully_connected(_6.get_output(0), 1024, w, b) _8 = network.add_activation(_7.get_output(0), trt.ActivationType.RELU)

w = np.random.rand(10,1024).astype(np.float32).reshape(-1) b = np.random.rand(10).astype(np.float32).reshape(-1) _9 = network.add_fully_connected(_8.get_output(0), 10, w, b) _10 = network.add_activation(_9.get_output(0), trt.ActivationType.RELU)

_11 = network.add_shuffle(_10.get_output(0)) _11.reshape_dims = [-1, 10]

_12 = network.add_softmax(_11.get_output(0)) _12.axes = 1 << 1

_13 = network.add_topk(_12.get_output(0), trt.TopKOperation.MAX, 1, 1 << 1)

network.mark_output(_13.get_output(1))

打印逐層信息

for i in range(network.num_layers): layer = network.get_layer(i) print(i,"%s,in=%d,out=%d,%s"%(str(layer.type)[10:],layer.num_inputs,layer.num_outputs,layer.name)) for j in range(layer.num_inputs): tensor =layer.get_input(j) if tensor == None: print("\tInput %2d:"%j,"None") else: print("\tInput %2d:%s,%s,%s"%(j,tensor.shape,str(tensor.dtype)[9:],tensor.name)) for j in range(layer.num_outputs): tensor =layer.get_output(j) if tensor == None: print("\tOutput %2d:"%j,"None") else: print("\tOutput %2d:%s,%s,%s"%(j,tensor.shape,str(tensor.dtype)[9:],tensor.name))

engineString = builder.build_serialized_network(network, config) if engineString == None: print("Failed building engine!") else: print("Succeeded building engine!") ```

  以上構建了一個網絡,並使用循環遍歷 network.layers,打印 layer type、layer num_inputs、layer num_outputs、layer name、layer input tensor、layer output tensor 等信息,層信息打印輸出如下:


  好了,以上分享了 TensorRT 中 Layer 和 Tensor 的區別。希望我的分享能對你的學習有一點幫助。


 【公眾號傳送】

《極智AI | TensorRT 中 Layer 和 Tensor 的區別》


logo_show.gif