极智AI | 讲解 TensorRT Constant 算子

  欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享

  大家好，我是极智视界，本文讲解一下 TensorRT Constant 算子。

  Constant 算子是指常量层，这个算子一般是什么时候使用呢：一般当下一个算子是两矩阵乘 或 两矩阵点乘 或 两矩阵拼接等这类两头输入的算子，而某一个矩阵需要离线读取时，就需要用到 Constant 算子来构建这个离线读取的张量。以上介绍了 Constant 算子一个使用场景，下面介绍 TensorRT 中 Constant 算子的具体怎么来添加。

  在 TensorRT 中如何构建一个 Constant 算子呢，来看：

# 通过 add_constant 添加 constant 算子
constantLayer = network.add_constant([1], np.array([1], dtype=np.float32))

# 重设常量数据
constantLayer.weights = data 

# 重设常量形状
constantLayer.shape = data.shape
1
2
3
4
5
6
7
8

  来看一个实际的例子：

import numpy as np
from cuda import cudart
import tensorrt as trt

# 输入张量 NCHW
nIn, cIn, hIn, wIn = 1, 3, 4, 5  # 输入张量 NCHW

# 输入数据
data = np.arange(nIn * cIn * hIn * wIn, dtype=np.float32).reshape(nIn, cIn, hIn, wIn) 

np.set_printoptions(precision=8, linewidth=200, suppress=True)
cudart.cudaDeviceSynchronize()

logger = trt.Logger(trt.Logger.ERROR)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
config = builder.create_builder_config()
#---------------------------------------------------------- --------------------# 替换部分
# 添加 constant 算子
constantLayer = network.add_constant(data.shape, data)
#---------------------------------------------------------- --------------------# 替换部分
network.mark_output(constantLayer.get_output(0))

engineString = builder.build_serialized_network(network, config)
engine = trt.Runtime(logger).deserialize_cuda_engine(engineString)
context = engine.create_execution_context()
_, stream = cudart.cudaStreamCreate()

outputH0 = np.empty(context.get_binding_shape(0), dtype=trt.nptype(engine.get_binding_dtype(0)))
_, outputD0 = cudart.cudaMallocAsync(outputH0.nbytes, stream)

context.execute_async_v2([int(outputD0)], stream)
cudart.cudaMemcpyAsync(outputH0.ctypes.data, outputD0, outputH0.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)
cudart.cudaStreamSynchronize(stream)

print("outputH0:", outputH0.shape)
print(outputH0)

cudart.cudaStreamDestroy(stream)
cudart.cudaFree(outputD0)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

• 输出张量形状 (1,3,4,5)
  [ [ [ 0. 1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19. ] [ 20. 21. 22. 23. 24. 25. 26. 27. 28. 29. 30. 31. 32. 33. 34. 35. 36. 37. 38. 39. ] [ 40. 41. 42. 43. 44. 45. 46. 47. 48. 49. 50. 51. 52. 53. 54. 55. 56. 57. 58. 59. ] ] ] \left[
  \begin{matrix} \left[\begin{matrix} \left[\begin{matrix} 0. & 1. & 2. & 3. & 4. \\ 5. & 6. & 7. & 8. & 9. \\ 10. & 11. & 12. & 13. & 14. \\ 15. & 16. & 17. & 18. & 19. \end{matrix}$$\text{begin{matrix} left[begin{matrix} left[begin{matrix} 0. \& 1. \& 2. \& 3. \& 4. 5. \& 6. \& 7. \& 8. \& 9. 10. \& 11. \& 12. \& 13. \& 14. 15. \& 16. \& 17. \& 18. \& 19. end{matrix}}$$
  \right] \left[
  20.25.30.35.21.26.31.36.22.27.32.37.23.28.33.38.24.29.34.39.$$\begin{array}{ccccc}20.& 21.& 22.& 23.& 24.\\ 25.& 26.& 27.& 28.& 29.\\ 30.& 31.& 32.& 33.& 34.\\ 35.& 36.& 37.& 38.& 39.\end{array}$$
  \right] \left[
  40.45.50.55.41.46.51.56.42.47.52.57.43.48.53.58.44.49.54.59.$$\begin{array}{ccccc}40.& 41.& 42.& 43.& 44.\\ 45.& 46.& 47.& 48.& 49.\\ 50.& 51.& 52.& 53.& 54.\\ 55.& 56.& 57.& 58.& 59.\end{array}$$
  \right] \end{matrix}\right] \end{matrix}\right] ⎣⎢⎢⎡​⎣⎢⎢⎡​⎣⎢⎢⎡​0.5.10.15.​1.6.11.16.​2.7.12.17.​3.8.13.18.​4.9.14.19.​⎦⎥⎥⎤​⎣⎢⎢⎡​20.25.30.35.​21.26.31.36.​22.27.32.37.​23.28.33.38.​24.29.34.39.​⎦⎥⎥⎤​⎣⎢⎢⎡​40.45.50.55.​41.46.51.56.​42.47.52.57.​43.48.53.58.​44.49.54.59.​⎦⎥⎥⎤​​⎦⎥⎥⎤​​⎦⎥⎥⎤​

  好了，以上分享了 讲解 TensorRT Constant 算子，希望我的分享能对你的学习有一点帮助。

 【公众号传送】

《极智AI | 讲解 TensorRT Constant 算子》

扫描下方二维码即可关注我的微信公众号【极智视界】，获取我的更多经验分享，让我们用极致+极客的心态来迎接AI ！