博客


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


大家好，我是极智视界，本文讲解一下 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)

：第一个参数为索引的 维度，第二个参数为索引的维度的序列号。

来看示例代码：

>>