极智AI | Attention 中 torch.chunk 的 TensorRT 实现

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

  大家好，我是极智视界，本文介绍一下 Attention 中 torch.chunk 的 TensorRT 实现。

  Attention 是 Transformer 中的重要模块，正如那句 Attention is all your need 。注意力机制最开始应用在自然语言处理领域，如机器翻译、语言理解，而后在 CV 领域、多模态领域也大放异彩，如 image caption 看图说话、ViT 等。这里拿 ViT 中的 Multi-Head Attention 为例，说说其中的生成 Q、K、V 的 TensorRT 实现。

1 ViT Multi-Head Attention

  先来说说 ViT 中的 Multi-Head Attention， 来源于论文《Attention Is All You Need》，示意如下：

  多头即将模型分为多个头，形成多个子空间，让模型去关注不同方面的信息，将 Scaled Dot-Product Attention 过程做 h 次，再把输出做 cat。这样做的目的是为了使网络能够综合利用多方面角度提取更加准确的表示，从而可以捕捉到更加丰富的特征，可以类比 CNN 中多个核分别提取特征的作用，原文是这么说的：

Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions.

在上面的介绍中，可以看到 Attention 中的最开始输入是 Q、K、V，那么这三个张量是怎么得到的呢？很简单，只有一句代码：

q, k, v = linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)

  其中有两个过程：

• linear
• chunk

  chunk 的作用就是对 linear 出来的张量进行分块，这里是等分成三块，以形成 Q、K、V。下面介绍 chunk 的 pytorch 和 TensorRT 实现。

2 torch.chunk 实现

  pytorch 中的 torch.chunk 方法可以对张量进行分块，方法定义如下：

torch.chunk(tensor, chunks, dim=0)

  其中的参数：

• tensor (Tensor)：待分块的输入张量；
• chunks (int)：分块的个数；
• dim (int)：沿着此维度进行分块；

  进行使用示例演示 (以下代码在终端)：

>>> import torch
>>> a = torch.randn(6, 4)
>>> a
tensor([[ 0.8268, -0.5770, -0.6231,  1.7666],
        [ 1.1914,  0.1997,  1.2418, -0.2340],
        [-0.1327, -0.0358, -0.4087, -0.7176],
        [ 1.3613, -1.8879, -0.6814, -0.2614],
        [-0.1122,  1.3339, -0.5488,  1.2073],
        [ 1.5515, -0.1518,  0.0484, -0.6843]])
>>> b = torch.chunk(a, 2, 1)
>>> b
(tensor([[ 0.8268, -0.5770],
        [ 1.1914,  0.1997],
        [-0.1327, -0.0358],
        [ 1.3613, -1.8879],
        [-0.1122,  1.3339],
        [ 1.5515, -0.1518]]), tensor([[-0.6231,  1.7666],
        [ 1.2418, -0.2340],
        [-0.4087, -0.7176],
        [-0.6814, -0.2614],
        [-0.5488,  1.2073],
        [ 0.0484, -0.6843]]))
>>> c = torch.chunk(a, 3, 0)
>>> c
(tensor([[ 0.8268, -0.5770, -0.6231,  1.7666],
        [ 1.1914,  0.1997,  1.2418, -0.2340]]), tensor([[-0.1327, -0.0358, -0.4087, -0.7176],
        [ 1.3613, -1.8879, -0.6814, -0.2614]]), tensor([[-0.1122,  1.3339, -0.5488,  1.2073],
        [ 1.5515, -0.1518,  0.0484, -0.6843]]))

3 TensorRT chunk 实现

  下面来说 chunk 算子的 TensorRT 实现，首先 TensorRT 中的标准算子并没有 chunk，所以这里我们需要做一些组合搭配，其核心主要用到了 Slice 算子，即切分，这里以四维张量为例：

for (int i = 0; i < chunks; ++i){
  nvinfer1::Dims start{4, 0, blockSize * i, 0, 0};
  nvinfer1::Dims size{4, bathSize, blockSize, 1, 1};
  NVINFER1::Dims stride{4, 1, 1, 1, 1};
  
  nvinfer1::ISliceLayer *Slice_layer = m_network->addSlice(front_Layer_out, start, size, stride);
  Layer[chunks_out[i]] = Slice_layer->getOutput(0);
}

  核心代码如上，可以看出 TensorRT 中实现起来并没有 pytorch 中那么灵活，其中的 blockSize 需要提前根据 chunks 和输入张量的维度计算好，且切分维度 dim 也需要在 start 和 size 中自己指定。

  好了，以上分享了 Attention 中 torch.chunk 的 TensorRT 实现。希望我的分享能对你的学习有一点帮助。

 【公众号传送】

《极智AI | Attention 中 torch.chunk 的 TensorRT 实现》

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