极智AI | 讲解 TensorRT 显式batch 和 隐式batch

  ​

​欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享​

​

  大家好，我是极智视界，本文讲解一下 TensorRT 显式batch 和 隐式batch。

  TensorRT 支持使用两种方式来指定网络的 Layout，即：​

​explicit batch​

​​ 和 ​

​implicit batch​

​。所谓 显式 和 隐式 的差别就在于 Batch 这一维，即 显式 ==> NCHW，隐式 ==> CHW。在比较老版本的 TensorRT 中一般就用 ​

​implicit batch​

​​ 隐式batch，而现在新的开始慢慢用 ​

​explicit batch​

​ 进行替代。这是为啥呢？这主要是因为，随着算法的不断发展，网络中新出现了很多新算子、新结构，有些时候需要 操控 batch 这个维度，这个时候如果使用 隐式batch 来开发，显然是不可行的。

  隐式batch 模式不能胜任的场景主要包括：

• Reducing across the batch dimension ==> batch 维度上的规约操作，如 [N, C, H, W] -> [1, C, H, W]；
• Reshaping the batch dimension ==> batch 维度上的变化，如 [N, C, H, W] -> [M, C, H, W]；
• Transposing the batch dimension with another dimension ==> 涉及 batch 的维度变换操作，如 [N, C, H, W] -> [C, H, W, N]；

  总的来说，就是涉及到 batch 的操作，​

​explicit batch​

​ 就无能为力了，而显而易见的是，隐式batch 的开发难度会低一些，因为开发自始至终，你只需要关注 CHW。有同学可能会关心的一点是，那么 我们如果要从 隐式batch模式 迁移到 显式batch模式 的工作量到底有多少呢。这个问题我觉得需要从两个方面进行回答：

• (1) 开发模式上如何设置成​

  ​explicit batch​

  ​ 模式；
• (2) 网络中算子 Layout 从 CHW -> NCHW，TensorRT 算子实现要改动的代码量；

  从第 (1) 个方面来说，很简单，只要在创建 ​

​INetworkDefinition​

​​ 的时候，使用 ​

​flag​

​ 来指定是 显式batch 还是 隐式batch。

  C++ 的做法是这样的，显式batch 时，手动指定 ​

​flag​

​​ 为 ​

​NetworkDefinitionCreationFlag::kEXPLICIT_BATCH​

​​ 以标志构建网络使用 ​

​explicit batch mode​

​​；若要使用 ​

​implicit batch mode​

​​ ，则 ​

​flag​

​​ 不进行手动指定 或 使用 ​

​createNetwork​

​ 方法来构建网络：

// cpp mark explicit batch mode
IBuilder* builder = ...;
INetworkDefinition* network = builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)))      

  Python 的代码是这样的，显式batch 时，手动指定 ​

​flag​

​​ 为 ​

​trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH​

​​ 以标志构建网络使用 ​

​explicit batch mode​

​​；若要使用 ​

​implicit batch mode​

​​ ，则 ​

​flag​

​ 不进行手动指定：

# python mark explicit batch mode
builder = trt.Builder(...)
builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))      

  从第 (2) 个方面来说，工作量其实是不太好评估的，这取决于你网络中涉及了多少 操控 batch 的操作，如果涉及的多，那改动自然就大。还有一点需要注意的是，​

​explicit batch​

​​ 中的 ​

​batch​

​​ 一般一直都是在 ​

​axis 0​

​​ 位置，除非特别指定 ​

​Transposing the batch dimension with another dimension​

​​ 可能会涉及 N 的 ​

​axis​

​​ 动一动，但这并不是常见的操作。因为一般一个高效的推理框架，会对你的 ​

​Data Layout​

​ 有严格要求，语义上就是 用 NCHW 还是 NHWC，框架是有一个期望的，而不是让你随意动，而反观 TensorRT，NCHW 的 Layout 是其期望。

  好了，以上分享了 讲解 TensorRT 显式batch 和 隐式batch，希望我的分享能对你的学习有一点帮助。