pooling2d

功能说明

通过不同的池化方式，对tensor_in上不同的滑动窗口进行信号的采样。

池化方式支持MAX、AVG、GMP、GAP四种：

• MAX：max pooling 是取每个滑动窗口覆盖的元素的最大值作为计算输出。
• AVG：avg pooling是取每个滑动窗口覆盖的元素的和除以窗口面积取平均值后作为输出。
• GMP：global max pooling 是 max pooling 的特殊模式，即 window 的大小和 feature map 的大小一致的 max pooling
• GAP：global avg pooling 是 avg pooling 的特殊模式，即 window 的大小和 feature map 的大小一致的 avg pooling

如下图，tensor_in按照pooling_mode为MAX、padding_mode为SAME时池化结果。

其中：

• input_w：tensor_in的W。
• input_h：tensor_in的H。
• kernel_w：window的W
• kernel_h：window的H
• pad_top：tensor_in的H方向顶部的pad行数，图示中此值为1
• pad_bottom：tensor_in的H方向底部的pad行数，图示中此值为1
• pad_left：tensor_in的W方向上左边的pad列数，图示中此值为1
• pad_right：tensor_in的W方向行右边的pad列数，图示中此值为1
• stride_w：stride的W
• stride_h：stride的H

您可以在“te/lang/cce/te_compute/pooling2d_compute.py”查看接口定义。

约束说明

此接口暂不支持与其他TBE DSL计算接口混合使用。

• 支持的输入数据类型：float16。
• 输入的tensor_in必须是五维，即符合5D-NC1HWC0格式排布的tensor。
• 输入的tensor_in的tensor的最后一维C0必须是16。
• window的维度必须为2，且为正整数，在[1, 32768]范围内。
• stride的维度必须为2且为正整数类型，stride的width、height需在[1, 63]范围内。
• 如果输入pad，那么pad的维度必须为4，pad值应 >= 0。
• dilation的维度为2且为正整数类型，在[1, 255]范围内。
• 当pooling_mode为MAX或AVG的VALID模式时需满足如下条件：

  out_w * window_h * window_w * C0 * SIZE_OF_FP16 + out_w * C0 * SIZE_OF_FP16 < ub_size

• 当pooling_mode为AVG的SAME模式时需满足如下条件：

  out_w * window_h * window_w * C0 * SIZE_OF_FP16 + out_w * C0 * SIZE_OF_FP16

  + out_w * C0 * SIZE_OF_FP16 < ub_size

• 当pooling_mode为AVG时，需满足 stride_h <= 2*window_h 且 stride_w <= 2*window_w。
• 当pooling_mode为AVG时，需满足 window的width*height < 256。
• 当pooling_mode为MAX时，需满足 window的width<=20 且 height <= 20
• 当pooling_mode为MAX或AVG时，tensor_in, pad, window 需满足：

  stride_h <= in_size_h + pad_top + pad_bottom – window_h

  stride_w <= in_size_w + pad_left + pad_right – window_w

• 当pooling_mode为GAP或GMP时，需满足 window_h = in_size_h 且 window_w = in_size_w。
• 当pooling_mode为GAP或GMP时，需满足 padding_mode = "VALID"。

• ub_size为UB可用大小。
• out_w为输出结果tensor的W。
• window_h为window的height。
• window_w为window的width。
• C0为tensor_in的C0。
• SIZE_OF_FP16为float16的大小。

函数原型

te.lang.cce.pooling2d(tensor_in, window, stride, pooling_mode, padding_mode="SAME", pad = (0,0,0,0), dilation = (1,1), data_mode=1, ceil_mode=0, fusion_params=None, impl_mode="high_performance")

参数说明

• tensor_in：输入的feature map，tvm.tensor类型。符合5D-NC1HWC0格式排布的tensor，各维度分别代表其NC1HWC0信息。
• window：输入的滑块大小信息，list、tuple类型。window[0]表示输入window的width，window[1]输入window的height。
• stride：输入的滑块移动步长信息，list、tuple类型。stride[0]表示window在feature map的W方向上移动的步长，stride[1]表示window在feature map的H方向上移动的步长。
• pooling_mode：pooling模式，支持“MAX”、“AVG"、“GMP”、“GAP”，分别表示max pooling, avg pooling, global max pooling, global avg pooling。
  • MAX：max pooling 是取每个滑动窗口覆盖的元素的最大值作为计算输出。
  • AVG：avg pooling是取每个滑动窗口覆盖的元素的和除以窗口面积取平均值后作为输出。
  • GMP：global max pooling 为 max pooling 的特殊模式，feature map 大小和 window 大小一致，取 feature map 元素的最大值作为计算输出。
  • GAP：global avg pooling 为 avg pooling 的特殊模式，feature map 大小和 window 大小一致，取 feature map 元素之和取均值后作为计算输出。
• padding_mode：padding模式，支持“VALID”、“SAME”，分别代表 不补pad、补pad。
  • VALID模式：

    当窗口在W方向或者H方向上的移动只能覆盖部分feature map时，丢弃掉这些不足一个完整窗口覆盖的数据，即：feature map 的这部分数据不参与计算。

    MAX , AVG, GMP, GAP 都包含VALID模式

  • SAME模式：

    当窗口在W方向或者H方向上的移动只能覆盖部分feature map时，补齐pad为0值至刚好能够满足一个完整窗口的覆盖，即：feature map 的这部分数据通过补 pad = 0 的方式参与计算。

    MAX and AVG pooling 包含SAME模式，GMP and GAP 没有SAME模式

• pad：补pad的数目，list、tuple类型。可选参数，用于兼容Caffe的pooling。pad[0], pad[1], pad[2], pad[3] 分别代表用户输入的在top, bottom, left, right方向补的pad，默认值为(0,0,0,0)。
• dilation：膨胀系数，list、tuple类型。可选参数，dilation[0]、dilation[1]分别表示window在H、W方向的膨胀系数，默认值为(1,1)。
• data_mode：模板使用类型，0: CAFFE_DATA_MODE, 1: TENSORFLOW_DATA_MODE。
• ceil_mode：Caffe模式下的round_mode 参数, 0:CEIL(default), 1:FLOOR。
• fusion_params：预留字段，开发者无需关注。
• impl_mode：算子运行时选择精度优先还是性能优先模式，有以下两种取值：
  • high_precision：表示精度优先，精度优先时，由于计算过程复杂，性能会有所下降。
  • high_performance：表示性能优先，性能优先时，精度会有所下降。

  默认值“high_performance”。

返回值

res_tensor：输出tensor，tvm.tensor类型，为符合5D-NC1HWC0格式排布的tensor。

将tensor_in 的shape信息记为[N, C1, H, W, C0=16]，window 的shape信息记为 [F, F]，stride 信息记为 [S, S]，则：

MAX and AVG 的VALID模式与SAME模式下输出tensor的shape信息计算方式分别如下所示：

• VALID模式下
  • N 和C 维度保持不变。
  • Hout 和 Wout 维度大小为：

    

• SAME模式下
  • N 和C 维度保持不变。
  • Hout 和Wout 维度大小为：

    

    其中，W为输入size，F为filter的size，S为步长，[]为向上取整符号。

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

调用示例

from tbe import tvm
import te.lang.cce
shape = (1, 2, 416, 416, 16) 
input_dtype = "float16"
data = tvm.placeholder(shape, name="data", dtype=input_dtype) 
res = te.lang.cce.pooling2d(data, (3, 3), (2, 2), "AVG", "SAME")
# res.shape = (1, 2, 208, 208, 16)