插件的数据输入格式

Data input format to Plugin

我有张量流模型,如图所示。并尝试转换为 TensorRT。

在转换为 TensorRT 时需要 upsample 节点的插件。

对于 Tensorflow 实现,假设输入大小为 1x3x4x19 形状并上采样为 1x12x14x19 张量。

在 TensorRT 插件中实现了同样的事情,考虑的流程如下。

tensorflow 中的 1x3x4x19 张量是

[[[...19channel data ...],[...19channel data ...],[...19channel data ...],[...19channel data ...]],

 [[...19channel data ...],[...19channel data ...],[...19channel data ...],[...19channel data ...]],

 [[...19channel data ...],[...19channel data ...],[...19channel data ...],[...19channel data ...]]]

首先扁平化为

[...19channel data ..., ...19channel data ..., ...19channel data ..., etc.,...19channel data ...]

展平后的数据长度为 228。

19 个通道的数据很难可视化。

所以再次使用3通道数据作为展平数据的例子。

[[[1,2,3],[4,5,6],[7,8,9],[10,11,12]],
 [[1,2,3],[4,5,6],[7,8,9],[10,11,12]],
 [[1,2,3],[4,5,6],[7,8,9],[10,11,12]]]

3 通道数据的展平数组是

[1,2,3,4,5,6,7,8,9,10,11,12,1,2,3,4,5,6,7,8,9,10,11,12,1,2,3,4,5,6,7,8,9,10,11,12]

扁平化的数据被输入到插件中进行上采样。 我在插件中的 CUDA 代码需要如上所示的扁平化数据。

但插件输出很奇怪,与 Tensorflow 的上采样输出不同。 我使用 Openpose1(如图所示)数据检查了 Tensorflow 的上采样操作。

插入TensorRT引擎的数据格式是否正确? 如果不是,如何将输入数据提供给插件?

插件输出看起来像输入被垂直压平了

[1,1,1,2,2,2,3,3,3...etc.]

经过一段时间,问题解决了。 插件输入是 NCHW 格式。 它就像 1,4,7,10,1,4,7,10,1,4,7,10,2,5,8,11,2,5,8,11,2,5,8,11 ,3,6,9,12,3,6,9,12,3,6,9,12.

因此,如果匹配 Tensorflow 的 NHWC 格式,则需要处理该数据并重新格式化回 NHWC 格式。