TensorRT权重顺序与Tensorflow [H，W，IN_C，OUT_C]比较

Question

我必须手动添加卷积层，因为我有一些Tensorflow解析器不支持的特殊操作。 TensorRT期望从.wts文件读取的权重顺序是什么？例如，Tensorflow的转换权重张量通常具有[H，W，IN_CHANNEL，OUT_CHANNEL]的顺序。我知道TensorRT期望输入数据按NCHW顺序排列，但是在将其写入.wts文件时，权重顺序也必须更改吗？如果是这样，TensorRT预期会得到什么订单？ [IN_CHANNEL，OUT_CHANNEL，H，W]？

Answer 1

快速摘要...如果您询问权重排序顺序，那么您可能还会担心输入数据的顺序。在这两个方面，这里发布的答案可能会为您提供大部分所需的东西：Run Tensorflow with NVIDIA TensorRT Inference Engine

其他详细信息...我最近使用自定义工具解决了这些问题，以下是我遇到的相关因素：

• 输入图像数据顺序，对于张量流为NHWC，对于张量为NCHW；并且在图像的通道内，颜色顺序例如RGB与BGR

• 权重按图层排序的顺序

  • 对于2D卷积，tensorflow使用RSCK（[filter_height，filter_width，input_depth，output_depth]），而tensorrt使用KCRS。

  • 对于2D卷积或池化层之后的密集层，调整权重排序顺序以使用不同的展平序列，有效地将张量流的RSCK转换为张量的KCRS，现在R和S代表整个输入层高度和宽度分别为C和以前一样，input_depth，现在输出深度K为致密层的神经元计数

  • 对于紧随致密层的致密层，将CK转换为KC阶

（注意：此答案假设您未在任何卷积中使用组）