pooling3d

功能说明

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

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

如下图所示,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

其中:

该接口支持pooling的基本功能, 但不支持出口量化功能。

函数原型

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)

参数说明

返回值

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信息计算方式分别如下所示:

约束说明

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

当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
  • ub_size为UB可用大小。
  • out_w为输出结果tensor的W。
  • window_h为window的height。
  • window_w为window的width。
  • C0为tensor_in的C0。
  • SIZE_OF_FP16为float16的大小。

支持的型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

Atlas 推理系列产品

Atlas A2训练系列产品

调用示例

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