pooling3d

功能说明

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

池化方式支持MAX和AVG两种：

• MAX：maxpooling3d是取每个滑动窗口覆盖的元素的最大值作为计算输出 。
• AVG：avgpooling3d是取每个滑动窗口覆盖的元素的平均值作为计算输出。

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

input_d = 4, input_h =4, input_w = 4

stride_d = 2, stride_h = 2, stride_w =2

kernel_d = 2, kernel_h = 2, kernel_w = 2

其中：

• input_d: tensor_in中的D
• input_h：tensor_in的H。
• input_w: tensor_in中的W
• kernel_d：window的D
• kernel_h：window的H
• kernel_w：window的W
• stride_d: strides中的D
• stride_h: strides中的H
• stride_w: strides中的W
• pad_top：tensor_in的D方向顶上的pad的行数，图示中此值为0
• pad_bottom：tensor_in的D方向底部的pad的行数，图示中此值为0
• pad_front: tensor_in的H方向前面的pad的行数，图示中此值为0
• pad_back: tensor_in的H方向后面的pad的行数，图示中此值为0
• pad_left：tensor_in的W方向上左边的pad列数，图示中此值为0
• pad_right：tensor_in的W方向行右边的pad列数，图示中此值为0

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

约束说明

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

• 支持的输入数据类型：float16。
• 输入的tensor_in必须是五维，即符合5D-NC1HWC0格式排布的tensor。
• 输入的tensor_in必须是NDC1HWC0格式。
• window的维度必须为3，且为正整数，在[1, 8]范围内。
• stride的维度必须为3且为正整数类型，stride的width、height需在[1, 8]范围内。
• 如果输入pad，那么pad的维度必须为6，pad值应 >= 0。
• dilation的维度为3且为正整数类型，此参数为接口保留，当前未产生效果。
• 当pooling_mode为MAX或AVG时，tensor_in, pad, window 需满足：

  stride_d <= in_size_d + pad_top + pad_bottom – window_d

  stride_h <= in_size_h + pad_front + pad_back – window_h

  stride_w <= in_size_w + pad_left + pad_right – window_w

函数原型

te.lang.cce.pooling3d(tensor_in, window, stride, padding_mode="SAME", pads=(0, 0, 0, 0, 0, 0), pooling_mode="MAX", dilation=(1, 1, 1), ceil_mode=0)

参数说明

• tensor_in：输入的feature map，tvm.tensor类型。符合6D-NDC1HWC0格式排布的tensor。
• window：输入的滑块大小信息，list、tuple类型。window[0]表示输入window的depth，window[1]表示输入window的width，window[2]输入window的height。
• stride：输入的滑块移动步长信息，list、tuple类型。stride[0]表示window在feature map的D方向上移动的步长，stride[1]表示window在feature map的W方向上移动的步长，stride[2]表示window在feature map的H方向上移动的步长。
• pooling_mode：pooling模式，支持“MAX”、“AVG"、分别表示max pooling3d, avg pooling3d,
• MAX：max pooling 是取每个滑动窗口覆盖的元素的最大值作为计算输出。
• AVG：avg pooling是取每个滑动窗口覆盖的元素的和除以窗口面积取平均值后作为输出。
• pads：补pad的数目，list、tuple类型。可选参数，用于兼容Caffe的pooling。pad[0], pad[1], pad[2], pad[3] ，pad[4], pad[5] 分别代表用户输入的在top, bottom, front, back, left, right方向补的pad，默认值为(0,0,0,0,0,0)。
• padding_mode：padding模式，支持“VALID”、“SAME”，分别代表 不补pad、补pad。
  • VALID模式：

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

  • SAME模式：

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

• dilation：膨胀系数，list、tuple类型。可选参数，dilation[0]、dilation[1]分别表示window在H、W方向的膨胀系数，默认值为(1,1,1), 当前该参数未生效。
• ceil_mode：Caffe模式下的round_mode 参数, 0:CEIL(default), 1:FLOOR。

返回值

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

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

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

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

    new_depth=new_height=new_width = CEIL(W-F+1/S)

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

    new_depth=new_height=new_width = CEIL(W/S)

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

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

调用示例

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