極智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，希望我的分享能對你的學習有一點幫助。