conv3d

功能说明

在给定6HD格式的Data和FracZ格式的Weight的情况下计算float16的3-D卷积。

接口可以支持bias。

Data tensor 的shape是6HD，即（N, D, C1, H, W, C0）；Weight Tensor 的shape是 FracZ，即 (KD*C1*KH*KW, Cout//C0_out, C0_out, C0)。

函数原型

conv3d(x, filter, filter_size, para_dict)

参数说明

• x：3d卷积的FeatureMap，tensor，6HD格式，目前支持float16类型。
• filter：3d卷积的Weight，tensor，FracZ格式，目前支持float16类型。
• filter_size：3d卷积的Weight。
• para_dict：字典格式，包含各种参数，后续参数扩展一般都在para_dict
  其中目前在para_dict里要传递的参数：

  • dsl_flag：是否支持UB自动融合，True：支持UB自动融合，False：不支持UB自动融合，默认值为False
  • bias_tensor：是否支持带bias，默认值为None
  • pads：list类型，3d卷积在FeatureMap DHW方向的补边。
  • strides：list类型，3d卷积在FeatureMap DHW方向的移动步长。
  • dilations：list类型，3d卷积Filter DHW方向的膨胀系数。
  • res_dtype：输出的数据类型。
  • mad_dtype：执行mad计算的数据类型。
  • kernel_name：算子名称。
  • group_dict：字典格式，包含支持3d group卷积相关的参数，通过下列计算过程获得。

    参数1：fmap_c，FeatureMap的C维度大小。

    参数2：cout，Weight的batch维度大小。

    参数3：groups，group卷积参数。

    参数4：cout0，为tbe_platform.C0_SIZE，默认值为16。

    参数5：cin0，为tbe_platform.C0_SIZE，默认值为16。

    具体计算公式：

    lcm(param1, param2)，计算最小公倍数。

    mag_factor0 = lcm(fmap_c // groups, cin0) // (fmap_c // groups)

    mag_factor1 = lcm(cout // groups, cout0) // (cout // groups)

    mag_factor = min(lcm(mag_factor0, mag_factor1), groups)

    cin1_g = (mag_factor * fmap_c // groups + cin0 - 1) // cin0

    cout_g = (mag_factor * cout // groups + cout0 - 1) // cout0 * cout0

    group_dict = {"real_g": (groups + mag_factor - 1) // mag_factor,

    "mag_factor": mag_factor,

    "cin1_g": cin1_g,

    "cout_g": cout_g,

    "cin_ori": fmap_c,

    "cout_ori": cout}

返回值

res_tensor：表示卷积计算的tensor，即卷积计算的结果输出。

约束说明

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

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

调用示例

from tbe import tvm
from tbe import dsl

shape_fmp_ndc1hwc0 = (1, 32, 1, 240, 352, 16)
fmp_dtype = "float16"
shape_filter = [16, 16, 3, 3, 3]
shape_w_frac_z = (27, 1, 16, 16)
w_dtype = "float16"

data = tvm.placeholder(shape_fmp_ndc1hwc0, name='Fmap', dtype=fmp_dtype)
weight = tvm.placeholder(shape_w_frac_z, name='Filter', dtype=w_dtype)

bias_tensor = None
pads = [1, 1, 1, 1, 1, 1]
stride_dhw = [1, 1, 1]
res_dtype = "float16"
mad_dtype = "float32"
kernel_name = "conv3d_1_32_240_352_1_3_3_3_1_16_1_1_1_SAME_NDHWC_1_0"
group_dict = {'real_g': 1, 'mag_factor': 1, 'cin1_g': 1, 'cout_g': 16, 'cin_ori': 1, 'cout_ori': 16}
dilation_dhw = [1, 1, 1]

para_dict = {
    "dsl_flag": False,
    "bias_tensor": bias_tensor,
    "pads": pads,
    "strides": stride_dhw,
    "res_dtype": res_dtype,
    "mad_dtype": mad_dtype,
    "kernel_name": kernel_name,
    "group_dict": group_dict,
    "dilations": dilation_dhw
}

conv_res = dsl.conv3d(data, weight, shape_filter, para_dict)