极智AI | 讲解 TensorRT Fully Connected 算子

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

  大家好,我是极智视界,本文讲解一下 TensorRT Fully Connected 算子。

  Fully Connected 也即 全连接层, 一般作为分类头或特征头使用。全连接层是个经典层,并不复杂,若没有偏置的话就是一个矩阵乘,如有偏置的话,就是一个矩阵乘然后接一个矩阵加。这里我们来看看 TensorRT 中 Fully Connected 的几种实现方式。

1 TensorRT 原生算子实现

  用 TensorRT Fully Connected 原生算子来实现肯定是最方便的,关键的几步如下:

  1. placeHolder = np.zeros(1, dtype=np.float32)
  2. # 添加全连接层
  3. fullyConnectedLayer = network.add_fully_connected(inputT0, 1, placeHolder, placeHolder)
  4. # 重设输出通道数
  5. fullyConnectedLayer.num_output_channels = cOut
  6. # 重设全连接权值
  7. fullyConnectedLayer.kernel = weight
  8. # 重设全连接偏置,bias 为可选参数,默认值 None
  9. fullyConnectedLayer.bias = bias

  来用一个完整的示例进行展示:

  1. import numpy as np
  2. from cuda import cudart
  3. import tensorrt as trt
  4. # 输入张量 NCHW
  5. nIn, cIn, hIn, wIn = 1, 3, 4, 5
  6. # 输出张量 C
  7. cOut = 2
  8. # 输入数据
  9. data = np.arange(cIn * hIn * wIn, dtype=np.float32).reshape(cIn, hIn, wIn)
  10. # 全连接权值
  11. weight = np.ones(cIn * hIn * wIn, dtype=np.float32)
  12. weight = np.concatenate([weight, -weight], 0).reshape(cOut, cIn, hIn, wIn)
  13. # 全连接偏置
  14. bias = np.zeros(cOut, dtype=np.float32)
  15. np.set_printoptions(precision=8, linewidth=200, suppress=True)
  16. cudart.cudaDeviceSynchronize()
  17. logger = trt.Logger(trt.Logger.ERROR)
  18. builder = trt.Builder(logger)
  19. network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
  20. config = builder.create_builder_config()
  21. inputT0 = network.add_input(inputT0, trt.DataType.FLOAT, (nIn, cIn, hIn, wIn))
  22. #-----------------------------------------------------------------------# 替换部分
  23. # 添加全连接层
  24. fullyConnectedLayer = network.add_fully_connected(inputT0, cOut, weight, bias)
  25. #-----------------------------------------------------------------------# 替换部分
  26. network.mark_output(fullyConnectedLayer.get_output(0))
  27. engineString = builder.build_serialized_network(network, config)
  28. engine = trt.Runtime(logger).deserialize_cuda_engine(engineString)
  29. context = engine.create_execution_context()
  30. _, stream = cudart.cudaStreamCreate()
  31. inputH0 = np.ascontiguousarray(data.reshape(-1))
  32. outputH0 = np.empty(context.get_binding_shape(1), dtype=trt.nptype(engine.get_binding_dtype(1)))
  33. _, inputD0 = cudart.cudaMallocAsync(inputH0.nbytes, stream)
  34. _, outputD0 = cudart.cudaMallocAsync(outputH0.nbytes, stream)
  35. cudart.cudaMemcpyAsync(inputD0, inputH0.ctypes.data, inputH0.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)
  36. context.execute_async_v2([int(inputD0), int(outputD0)], stream)
  37. cudart.cudaMemcpyAsync(outputH0.ctypes.data, outputD0, outputH0.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)
  38. cudart.cudaStreamSynchronize(stream)
  39. print("inputH0 :", data.shape)
  40. print(data)
  41. print("outputH0:", outputH0.shape)
  42. print(outputH0)
  43. cudart.cudaStreamDestroy(stream)
  44. cudart.cudaFree(inputD0)
  45. cudart.cudaFree(outputD0)
  • 输入张量形状 (1,3,4,5)
    $$
    \left[\begin{matrix}
    \left[\begin{matrix}

    1. \left[\begin{matrix}
    2. 0. & 1. & 2. & 3. & 4. \\
    3. 5. & 6. & 7. & 8. & 9. \\
    4. 10. & 11. & 12. & 13. & 14. \\
    5. 15. & 16. & 17. & 18. & 19.
    6. \end{matrix}\right]
    7. \left[\begin{matrix}
    8. 20. & 21. & 22. & 23. & 24. \\
    9. 25. & 26. & 27. & 28. & 29. \\
    10. 30. & 31. & 32. & 33. & 34. \\
    11. 35. & 36. & 37. & 38. & 39.
    12. \end{matrix}\right]
    13. \left[\begin{matrix}
    14. 40. & 41. & 42. & 43. & 44. \\
    15. 45. & 46. & 47. & 48. & 49. \\
    16. 50. & 51. & 52. & 53. & 54. \\
    17. 55. & 56. & 57. & 58. & 59.
    18. \end{matrix}\right]

    \end{matrix}\right]
    \end{matrix}\right]
    $$

  • 输出张量形状 (1,2,1,1)
    $$
    \left[\begin{matrix}
    \left[\begin{matrix}

    1. \left[\begin{matrix}
    2. \left[\begin{matrix}
    3. 1770.
    4. \end{matrix}\right]
    5. \end{matrix}\right] \\
    6. \left[\begin{matrix}
    7. \left[\begin{matrix}
    8. -1770.
    9. \end{matrix}\right]
    10. \end{matrix}\right]

    \end{matrix}\right]
    \end{matrix}\right]
    $$

  • 计算过程:$output = X \cdot W^{T} + bias$
    $$
    \begin{aligned}
    &= X.reshape(nIn,cIn \cdot hIn \cdot wIn) * W.reshape(cOut,cIn \cdot hIn \cdot wIn).transpose() \
    &= \left[\begin{matrix} 0 & 1 & 2 & \cdots & 59 \end{matrix}\right] +

    1. \left[\begin{matrix} 1 & -1 \\ 1 & -1 \\ \cdots & \cdots \\ 1 & -1 \end{matrix}\right] +
    2. \left[\begin{matrix} 0 & 0 \end{matrix}\right] \\

    &= \left[\begin{matrix} 1770 & -1770 \end{matrix}\right]
    \end{aligned}
    $$

  • Dynamic Shape 模式下,最低 3 维尺寸必须是构建期常量,不可为 -1

2 TensorRT 矩阵乘加实现

  然而全连接层又可以看成 一个矩阵乘接一个矩阵加。来看怎么做的:

  1. # 矩阵乘
  2. factorShape0 = weight.shape
  3. constantLayer0 = network.add_constant(factorShape0, np.ones(factorShape0factorShape0 = data.shape, dtype=np.float32))
  4. matrixMultiplyLayer = network.add_matrix_multiply(inputT0, trt.MatrixOperation.NONE, constantLayer0.get_output(0), trt.MatrixOperation.NONE)
  5. matrixMultiplyLayer.op0 = trt.MatrixOperation.NONE
  6. matrixMultiplyLayer.op1 = trt.MatrixOperation.TRANSPOSE
  7. # 矩阵加 (偏置)
  8. factorShape1 = bias.shape
  9. constantLayer1 = network.add_constant(factorShape1, np.ones(factorShape1, dtype=np.float32))
  10. biasLayer = network.add_elementwise(matrixMultiplyLayer.get_output(0), constantLayer1.get_output(0), trt.ElementWiseOperation.SUM)
  11. # get output
  12. output = biasLayer.get_output(0)

  这样就用 TensorRT 的乘加实现了 Fully Connected 算子。

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


 【公众号传送】


logo_show.gif


文章标签:

原文连接:https://juejin.cn/post/7116517454699823141

相关推荐

『时代』杂志:元宇宙将改变世界;健身教练:AI让我丢工作!有话说北欧人工智能夏令营资料大公开;深度学习书籍TOP5 | ShowMeAI资讯日报

图像文本跨模态细粒度语义对齐-置信度校正机制 AAAI2022

【Kaggle】如何有效避免OOM(out of memory)和漫长的炼丹过程

手把手教你运行YOLOv6(超详细)

OneFlow v0.8.0正式发布

transformer一统天下?depth-wise conv有话要说

OREPA:阿里提出训练也很快的重参数策略,内存减半,速度加倍 | CVPR 2022

【毕业设计】深度学习 opencv python 实现中国交通标志识别

【论文阅读|深读】RDAA:Role Discovery-Guided Network Embedding Based on Autoencoder and A

从“AI玩具”到“创作工具”的云原生改造之路

【深度学习】(一)机器学习基础

Keras深度学习实战(15)——从零开始实现YOLO目标检测

LabelImg(目标检测标注工具)的安装与使用教程

DeiT:注意力也能蒸馏

做一件有趣的事,尝试学着自己动手写一个深度学习框架(1)—深入反向传播

科学计算库Numpy基础&提升(理解+重要函数讲解)

极智AI | 变形金刚大家族 Transformer ViT CLIP BLIP BERT 模型结构

【大规模训练】transformer 中的张量模型并行

深度学习的训练、预测过程详解【以LeNet模型和CIFAR10数据集为例】

滚蛋吧小广告!我现在用命令行解压缩;当哥白尼遇上人工智能;一份傲娇的深度学习技术清单;一个视频尽览旷视20项前沿技术 | ShowMeAI资讯日报