极智AI | 讲解 TensorRT 怎么实现 torch.select 层

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

  大家好，我是极智视界，本文讲解一下 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)：第一个参数为索引的 维度，第二个参数为索引的维度的序列号。

  来看示例代码：

>>> 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])
1
2
3
4
5
6
7
8

2 TensorRT 实现 torch.select 层

  分析一下：在上面的 pytorch 介绍和示例演示中可以看出，select 可以通过 类似切片的操作 + 取我们想要的数据 来完成，自然在做 TensorRT 的实现的时候也可以往这个思路走。torch.select 主要由 dim 、index 两个因子来控制 取 数据的粒度，而TensorRT 在用 Slice 去切 Tensor 的时候，一般由 start、size、stride 三个因子来控制 切 的粒度。这样，其实 select 的 dim 和 index 完全可以转换为 Slice 的三个因子去控制。

  来用代码进行讲解：

/// 以下是 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
1
2
3
4
5
6
7
8
9
10
11

  以上 TensorRT slice <==> torch select 因子对照起来说明，应该会比较清楚一些。

  好了，以上分享了 讲解 TensorRT 怎么实现 torch.select 层。希望我的分享能对你的学习有一点帮助。

 【公众号传送】

《极智AI | 讲解 TensorRT 怎么实现 torch.select 层》

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