TensorRT 的上采样张量

Upsampling Tensor for TensorRT

tensorflow模型转为TensorRT,Tensorflow的ResizeArea(图中upsample)需要实现plugin。

所以ResizeArea是在CUDA中实现的。 我的 TensorRT 输入是 NCHW 格式。

uff_path = model_to_uff(model_path)
parser.register_input(ModelData.INPUT_NAME, (3, height, width), trt.UffInputOrder.NCHW)
parser.register_output(ModelData.OUTPUT_NAME)
parser.parse(uff_path, network)

所以我的 CUDA 代码实现了 NCHW 重采样。 我想确保我的重采样格式是正确的。

Method_1

NCHW 调整区域大小(4 倍上采样)样本。

channel_1      channel_2        channel_3
3,1,2,0,       0,4,3,1,         2,0,2,3,
3,0,1,2,       0,1,2,1,         2,0,4,2,
4,1,2,2,       1,3,2,4,         2,3,4,2,

channel_1                         channel_2                         channel_3  
3,3,3,3,1,1,1,1,2,2,2,2,0,0,0,0,  0,0,0,0,4,4,4,4,3,3,3,3,1,1,1,1,  2,2,2,2,0,0,0,0,2,2,2,2,3,3,3,3,
3,3,3,3,1,1,1,1,2,2,2,2,0,0,0,0,  0,0,0,0,4,4,4,4,3,3,3,3,1,1,1,1,  2,2,2,2,0,0,0,0,2,2,2,2,3,3,3,3,
3,3,3,3,1,1,1,1,2,2,2,2,0,0,0,0,  0,0,0,0,4,4,4,4,3,3,3,3,1,1,1,1,  2,2,2,2,0,0,0,0,2,2,2,2,3,3,3,3,
3,3,3,3,1,1,1,1,2,2,2,2,0,0,0,0,  0,0,0,0,4,4,4,4,3,3,3,3,1,1,1,1,  2,2,2,2,0,0,0,0,2,2,2,2,3,3,3,3,
3,3,3,3,0,0,0,0,1,1,1,1,2,2,2,2,  0,0,0,0,1,1,1,1,2,2,2,2,1,1,1,1,  2,2,2,2,0,0,0,0,4,4,4,4,2,2,2,2,
3,3,3,3,0,0,0,0,1,1,1,1,2,2,2,2,  0,0,0,0,1,1,1,1,2,2,2,2,1,1,1,1,  2,2,2,2,0,0,0,0,4,4,4,4,2,2,2,2,
3,3,3,3,0,0,0,0,1,1,1,1,2,2,2,2,  0,0,0,0,1,1,1,1,2,2,2,2,1,1,1,1,  2,2,2,2,0,0,0,0,4,4,4,4,2,2,2,2,
3,3,3,3,0,0,0,0,1,1,1,1,2,2,2,2,  0,0,0,0,1,1,1,1,2,2,2,2,1,1,1,1,  2,2,2,2,0,0,0,0,4,4,4,4,2,2,2,2,
4,4,4,4,1,1,1,1,2,2,2,2,2,2,2,2,  1,1,1,1,3,3,3,3,2,2,2,2,4,4,4,4,  2,2,2,2,3,3,3,3,4,4,4,4,2,2,2,2,
4,4,4,4,1,1,1,1,2,2,2,2,2,2,2,2,  1,1,1,1,3,3,3,3,2,2,2,2,4,4,4,4,  2,2,2,2,3,3,3,3,4,4,4,4,2,2,2,2,
4,4,4,4,1,1,1,1,2,2,2,2,2,2,2,2,  1,1,1,1,3,3,3,3,2,2,2,2,4,4,4,4,  2,2,2,2,3,3,3,3,4,4,4,4,2,2,2,2,
4,4,4,4,1,1,1,1,2,2,2,2,2,2,2,2,  1,1,1,1,3,3,3,3,2,2,2,2,4,4,4,4,  2,2,2,2,3,3,3,3,4,4,4,4,2,2,2,2,

每个像素被上采样4次(例如,第一个像素3被水平和垂直上采样4次)。这被认为是 NCHW 格式上采样。

Method_2

另一种实施方式被认为是 NHWC 格式的上采样。

3 通道数据 (83,86,77) 水平和垂直上采样。

Method_1 是 NCHW 上采样的正确方法吗?

看来方法1还可以,因为TensorRT需要CHW,如the docs中所述,NHWC是TF格式,你是否考虑在你的插件层中使用AlignCorners?另请注意,调整大小是最近邻,在 pytorch 中我使用 onnx-trt 进行双线性插值,这给出了更好的结果(在分割的情况下,也许对于你的情况 nn 是可以的)。

花了一些时间。问题已解决。 TensorRT 以 NCHW 格式工作。 Tensorflow 模型采用 NHWC 格式。所以在我的插件中,需要以 NCHW 格式进行上采样,但输出需要更改为 NHWC 格式,以便它可以连接到下一个 Tensorflow 操作。