極智AI | 講解 TensorRT 怎麼實現 torch.select 層

語言: CN / TW / HK

歡迎關注我的公眾號 [極智視界],獲取我的更多筆記分享

  大家好,我是極智視界,本文講解一下 TensorRT 實現 torch.select 層。

  我們知道,有很多算子都不會在 TensorRT 的原生支持算子列表裏,當然這裏要講的 select 算子也是一樣。然而,大部分的算子通過一些 TensorRT 原生算子的再組合就能實現,像 select、hardswich 等算子都可以這麼去做,但像 layernorm 等其他一些不容易通過原生算子組合實現的,直接用 plugin 實現會方便一些。

1 torch.select 介紹

torch.select 類似切片操作,如 切片 x[:, 0, :] 等價於 x.select(dim=1, index=0)

  其中 select(dim, index):第一個參數為索引的 維度,第二個參數為索引的維度的序列號。

  來看示例代碼:

```python

import torch a = torch.randn((3, 4)) a tensor([[-2.2622, 0.9470, -1.5170, -1.2614], [ 1.7269, 0.7789, 2.0953, -1.1928], [ 0.6136, -1.3214, 0.7611, -0.9582]]) a.select(dim=1, index=1) # 取第1個維度中索引為1的值 tensor([ 0.9470, 0.7789, -1.3214]) ```

2 TensorRT 實現 torch.select 層

  分析一下:在上面的 pytorch 介紹和示例演示中可以看出,select 可以通過 類似切片的操作 + 取我們想要的數據 來完成,自然在做 TensorRT 的實現的時候也可以往這個思路走。torch.select 主要由 dimindex 兩個因子來控制 數據的粒度,而TensorRT 在用 Slice 去切 Tensor 的時候,一般由 startsizestride 三個因子來控制 的粒度。這樣,其實 select 的 dimindex 完全可以轉換為 Slice 的三個因子去控制。

  來用代碼進行講解:

```c++ /// 以下是 explicit 模式的寫法,explicit 模式需要考慮 batch, 所以是四維的 // 假設輸入input shape 為 [N, C, H, W] ==> [32, 50, 1, 512] nvinfer1::Dims start{ 4, 0, 0, 0, 0 }; nvinfer1::Dims size{ 4, 32, 1, 1, 512 }; // 這裏相當於select 中 dim = 1, 因為是取了 [x, 50, x, x] 50 這個維度 nvinfer1::Dims stride{ 4, 1, 1, 1, 1 };

// 添加 Slice 層 nvinfer1::ISliceLayer slice = m_network->addSlice(input, start, size, stride); // 取數據 auto output = slice->getOutput(0); // 這裏相當於 select 中 index = 0 // auto output = slice->getOutput(1); // 這裏相當於 select 中 index = 1 ```

  以上 TensorRT slice <==> torch select 因子對照起來説明,應該會比較清楚一些。

  好了,以上分享了 講解 TensorRT 怎麼實現 torch.select 層。希望我的分享能對你的學習有一點幫助。


 【公眾號傳送】

《極智AI | 講解 TensorRT 怎麼實現 torch.select 層》


logo_show.gif