欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享
Olá a todos, sou Jizhi Vision. Este artigo explica como o TensorRT implementa a camada torch.select.
Sabemos que existem muitos operadores que não estão na lista de operadores suportados nativamente pelo TensorRT e, claro, o operador select a ser discutido aqui é o mesmo. No entanto, a maioria dos operadores pode ser implementada recombinando alguns operadores nativos do TensorRT. Operadores como select e hardswich podem fazer isso, mas alguns outros operadores, como layernorm, não são fáceis de serem implementados combinando operadores nativos. Será mais conveniente usar plugin alcançar.
1 tocha.selecione introdução
torch.selectSemelhante às operações de fatiamento, como fatiar é x[:, 0, :]equivalente a x.select(dim=1, index=0).
Onde select(dim, index): o primeiro parâmetro é a dimensão do índice e o segundo parâmetro é o número de série da dimensão do índice.
Dê uma olhada no código de exemplo:
>>> 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 implementa a camada torch.select
Análise: Como pode ser visto na introdução do pytorch acima e na demonstração do exemplo, a seleção pode ser feita através de operações semelhantes a fatiar + buscar os dados que desejamos. Naturalmente, também podemos seguir essa ideia ao implementar o TensorRT. torch.select controla principalmente a granularidade dos dados obtidos por dimdois fatoresindex , e o TensorRT geralmente controla a granularidade do corte por três fatores ao usar Slice para cortar tensores . Desta forma, de fato, e pode ser totalmente convertido em três fatores de Slice para controlar.startsizestridedimindex
Vamos explicar com código:
/// 以下是 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 层。希望我的分享能对你的学习有一点帮助。
【公众号传送】