tensorflow - TensorRT 权重顺序与 Tensorflow [H,W,IN_C,OUT_C] 的比较

Question

我必须手动添加一个卷积层，因为我有一些 Tensorflow 解析器不支持的特殊操作。TensorRT 期望从 .wts 文件中读取的权重顺序是什么？例如，Tensorflow 的 conv2d 权重张量通常具有 [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

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

• 输入图像数据顺序，tensorflow为NHWC，tensorrt为NCHW；在图像的通道内，颜色顺序，例如 RGB vs BGR

• 权重按层排序

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

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

  • 对于密集层之后的密集层，将 CK 转换为 KC 顺序

（注意：这个答案假设您没有在任何卷积中使用组）