极智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)：第一个参数为索引的 维度，第二个参数为索引的维度的序列号。

  来看示例代码：

>>