非连续转连续优化

原理

• PyTorch共享Storage机制。

  PyTorch中的tensor对象由表示层和存储层（Storage）构成，表示层主要包含tensor的形状、步长、类型和是否连续等信息，存储层为连续内存的一维数组。大多数情况下，每个tensor都有独立的表示层和存储层，但通过对tensor进行View类操作后，原始tensor和转换后的tensor表示层信息不同，但实际共享同一个存储层。

• PyTorch中tensor存储与读取原则是行优先。
  行优先是多维数组以一维展开的一种方式，举例如下：

  >>>t = torch.arange(12).reshape(3,4)
  >>>t
  tensor([[ 0,  1,  2,  3],        
          [ 4,  5,  6,  7],        
          [ 8,  9, 10, 11]])
  >>>t.stride()
  (4, 1)
  >>>t.is_contiguous()
  True

  这里定义了一个二维数组t，t的逻辑结构如下。

  图1 逻辑结构
  
  数组t在内存中实际以一维数组形式存储，可通过flatten查看。

  >>> t.flatten()
  tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])

  其存储物理结构如下：

  图2 一维展开
  

  t.stride()表示在指定维度dim中从一个元素跳到下一个元素所需的步长，t有0和1两个dim。沿着dim0（即纵向），从一个元素跳到下一个元素（如从0到4）要经过1、2、3、4，四个元素；沿着dim1（即横向），从一个元素跳到下一个元素（如从1到2）只经过2，一个元素。因此t.stride()为(4, 1)。

• PyTorch中的非连续tensor。
  tensor经过View非连续类（如transpose、permute等）操作后，转换后的tensor表示层信息不同但存储层相同。按照行优先展开时，存在数学上相邻的元素在内存上不再连续排布的情况。如对上述t进行transpose操作后结果如下。

  >>> t1 = t.transpose(0, 1)
  >>> t1
  tensor([[ 0,  4,  8],
          [ 1,  5,  9],
          [ 2,  6, 10],
          [ 3,  7, 11]])
  >>> t1.stride()
  (1, 4)
  >>> t1.is_contiguous()
  False
  >>> t1.flatten()
  tensor([ 0,  4,  8,  1,  5,  9,  2,  6, 10,  3,  7, 11])
  >>> t.data_ptr() == t1.data_ptr()
  True

  t1的逻辑结构如下，通过data_ptr方法可知t和t1首元素地址相同，可以判断为共用Storage。

  图3 t1逻辑结构
  
• 非连续转连续。

  NPU采用SIMD指令架构，对访问内存有连续性要求，需要将非连续的tensor转为连续的tensor。PyTorch中由非连续转连续方法为contiguous，例如：

  >>> t2 = t1.contiguous()
  >>> t2
  tensor([[ 0,  4,  8],
          [ 1,  5,  9],
          [ 2,  6, 10],
          [ 3,  7, 11]])
  >>> t2.data_ptr() == t1.data_ptr()
  False

  由以上可知，非连续转连续会重新开辟一块内存用来存储转换后的tensor，此过程存在性能开销。

问题定位

在模型或训练脚本中使用了View非连续类操作，如调用了transpose、narrow、select、permute、chunk、split等框架类算子，框架会调用format_contiguous函数对其进行校验，生成一个匹配且连续的tensor。

1. 参见Profiling数据采集及分析在训练脚本中使能profiling。
2. 通过查看Profiling数据，如发现在非连续转连续时耗时明显，可定位至对应算子，例如图4所示，format_contigousV2为转连续算子。
   图4 大量转连续算子
   

优化策略

• 使用计算类算子代替View类算子，view类算子指能够进行维度变换的算子，例如：view()、transpose()、permute()。

  使用index_select代替torch.transpose(x, 1, 2).contiguous。

  修改前，完整原始操作可参见开源链接：

  # 原始channel_shuffle操作
  def channel_shuffle(x, groups):
      # type: (torch.Tensor, int) -> torch.Tensor
      batchsize, num_channels, height, width = x.data.size()
      channels_per_group = num_channels // groups    # reshape
      x = x.view(batchsize, groups,
                 channels_per_group, height, width)
  
      x = torch.transpose(x, 1, 2).contiguous()
  
      # flatten
      x = x.view(batchsize, -1, height, width)
  
      return x

  同等语义修改：

  def channel_shuffle_index_select(x, groups=2):
      N, C, H, W = x.shape
      inp = C
  
      # channel_shuffle操作是对C维按一定规则的重排的工作，可以被表达为一次简单的重排
      group_len = inp // groups
      index = torch.from_numpy(np.array(list(range(inp))).reshape(groups, group_len).transpose(1, 0).flatten()).long()
  
      x = x.index_select(1, index)
      return x

• 如果存在对同一tensor的多次非连续操作，则可通过优先将其转连续以避免多次转换。