大家好，我是极智视界，本文讲解一下 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 Constant 算子，希望我的分享能对你的学习有一点帮助。
