TensorRT常用重点算子实现与原理讲解

​
目录

TensorRT Fully Connected 算子

1.TensorRT 原生算子实现

2.TensorRT 矩阵乘加实现

TensorRT Constant 算子

TensorRT 怎么实现 torch.select 层

1.torch.select 介绍

2.TensorRT 实现 torch.select 层

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

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

placeHolder = np.zeros(1, dtype=np.float32)

添加全连接层

fullyConnectedLayer = network.add_fully_connected(inputT0, 1, placeHolder, placeHolder)

重设输出通道数

fullyConnectedLayer.num_output_channels = cOut

重设全连接权值

fullyConnectedLayer.kernel = weight

重设全连接偏置，bias 为可选参数，默认值 None

fullyConnectedLayer.bias = bias
  来用一个完整的示例进行展示：

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

输入张量 NCHW

nIn, cIn, hIn, wIn = 1, 3, 4, 5

输出张量 C

cOut = 2

输入数据

data = np.arange(cIn * hIn * wIn, dtype=np.float32).reshape(cIn, hIn, wIn)

全连接权值

weight = np.ones(cIn * hIn * wIn, dtype=np.float32)
weight = np.concatenate([weight, -weight], 0).reshape(cOut, cIn, hIn, wIn)

全连接偏置

bias = np.zeros(cOut, dtype=np.float32)

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()
inputT0 = network.add_input(‘inputT0’, trt.DataType.FLOAT, (nIn, cIn, hIn, wIn))
#-----------------------------------------------------------------------# 替换部分

添加全连接层

fullyConnectedLayer = network.add_fully_connected(inputT0, cOut, weight, bias)
#-----------------------------------------------------------------------# 替换部分
network.mark_output(fullyConnectedLayer.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()

inputH0 = np.ascontiguousarray(data.reshape(-1))
outputH0 = np.empty(context.get_binding_shape(1), dtype=trt.nptype(engine.get_binding_dtype(1)))
_, inputD0 = cudart.cudaMallocAsync(inputH0.nbytes, stream)
_, outputD0 = cudart.cudaMallocAsync(outputH0.nbytes, stream)

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

print(“inputH0 :”, data.shape)
print(data)
print(“outputH0:”, outputH0.shape)
print(outputH0)

cudart.cudaStreamDestroy(stream)
cudart.cudaFree(inputD0)
cudart.cudaFree(outputD0)
输入张量形状 (1,3,4,5)
输出张量形状 (1,2,1,1)
计算过程：

2.TensorRT 矩阵乘加实现
  然而全连接层又可以看成 一个矩阵乘接一个矩阵加。来看怎么做的：

矩阵乘

factorShape0 = weight.shape
constantLayer0 = network.add_constant(factorShape0, np.ones(factorShape0factorShape0 = data.shape, dtype=np.float32))
matrixMultiplyLayer = network.add_matrix_multiply(inputT0, trt.MatrixOperation.NONE, constantLayer0.get_output(0), trt.MatrixOperation.NONE)
matrixMultiplyLayer.op0 = trt.MatrixOperation.NONE
matrixMultiplyLayer.op1 = trt.MatrixOperation.TRANSPOSE

矩阵加 (偏置)

factorShape1 = bias.shape
constantLayer1 = network.add_constant(factorShape1, np.ones(factorShape1, dtype=np.float32))
biasLayer = network.add_elementwise(matrixMultiplyLayer.get_output(0), constantLayer1.get_output(0), trt.ElementWiseOperation.SUM)

get output

output = biasLayer.get_output(0)
 这样就用 TensorRT 的乘加实现了 Fully Connected 算子。

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
来看一个实际的例子：

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,3,4,5)
 
TensorRT 怎么实现 torch.select 层
我们知道，有很多算子都不会在 TensorRT 的原生支持算子列表里，当然这里要讲的 select 算子也是一样。然而，大部分的算子通过一些 TensorRT 原生算子的再组合就能实现，像 select、hardswich 等算子都可以这么去做，但像 layernorm 等其他一些不容易通过原生算子组合实现的，直接用 plugin 实现会方便一些。

1.torch.select 介绍
  torch.select 类似切片操作，如 切片 x[:, 0, :] 等价于 x.select(dim=1, index=0)。 其中 select(dim, index)：第一个参数为索引的 维度，第二个参数为索引的维度的序列号。来看示例代码：

import torch
a = torch.randn((3, 4))
a
tensor([[-2.2622, 0.9470, -1.5170, -1.2614],
[ 1.7269, 0.7789, 2.0953, -1.1928],
[ 0.6136, -1.3214, 0.7611, -0.9582]])
a.select(dim=1, index=1) # 取第1个维度中索引为1的值
tensor([ 0.9470, 0.7789, -1.3214])
2.TensorRT 实现 torch.select 层
  分析一下：在上面的 pytorch 介绍和示例演示中可以看出，select 可以通过 类似切片的操作 + 取我们想要的数据 来完成，自然在做 TensorRT 的实现的时候也可以往这个思路走。torch.select 主要由 dim 、index 两个因子来控制取数据的粒度，而TensorRT 在用 Slice 去切 Tensor 的时候，一般由 start、size、stride 三个因子来控制切的粒度。这样，其实 select 的 dim 和 index 完全可以转换为 Slice 的三个因子去控制。

来用代码进行讲解：

/// 以下是 explicit 模式的写法，explicit 模式需要考虑 batch, 所以是四维的
// 假设输入input shape 为 [N, C, H, W] ==> [32, 50, 1, 512]
nvinfer1::Dims start{ 4, 0, 0, 0, 0 };
nvinfer1::Dims size{ 4, 32, 1, 1, 512 }; // 这里相当于select 中 dim = 1, 因为是取了 [x, 50, x, x] 50 这个维度
nvinfer1::Dims stride{ 4, 1, 1, 1, 1 };

// 添加 Slice 层
nvinfer1::ISliceLayer *slice = m_network->addSlice(*input, start, size, stride);
// 取数据
auto output = slice->getOutput(0); // 这里相当于 select 中 index = 0
// auto output = slice->getOutput(1); // 这里相当于 select 中 index = 1
 以上 TensorRT slice <==> torch select 因子对照起来说明，应该会比较清楚一些。

好了，以上分享了TensorRT常用重点算子实现与原理讲解。希望我的分享能对你的学习有一点帮助。

​