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

Posted 2022-08-06 极智视界

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

  大家好，我是极智视界，本文讲解一下 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，希望我的分享能对你的学习有一点帮助。

 【公众号传送】

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

扫描下方二维码即可关注我的微信公众号【极智视界】，获取我的更多经验分享，让我们用极致+极客的心态来迎接AI ！