欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享
大家好，我是极智视界，本文讲解一下 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])
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
以上 TensorRT slice <==> torch select 因子对照起来说明，应该会比较清楚一些。
好了，以上分享了 讲解 TensorRT 怎么实现 torch.select 层。希望我的分享能对你的学习有一点帮助。
【公众号传送】