conv2d

功能说明

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

Data tensor 的shape是5HD,即(N, C1, H, W, C0);Weight Tensor 的shape是 FracZ,即 (C1*KH*KW, Cout//C0_out, C0_out, C0)。

接口可以支持bias。

函数原型

conv(data, weight, para_dict, optim_dict=None, dsl_flag=True)

参数说明

返回值

res_tensor:表示卷积计算的tensor,即卷积计算的结果输出。

约束说明

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

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

调用示例

对于Batch=1,C=32,H=16,W=8的FeatureMap、Cout=32,KernelH=KernelW=1的Weight的float16的2D卷积,

对应的5HD的FeatureMap tensor的shape为(Batch,C/16,H,W,16)=(1,2,16,8,16)

对应的FracZ的Weight tensor的shape为 (KernelH*KernelW*C/16, Cout/16, 16, 16)=(2,2,16,16)

对应的Bias tensor的shape为(Cout,)=(32,)

from tbe import tvm
from tbe import dsl
 
shape_in = (1, 2, 16, 8, 16)
shape_w = (2, 2, 16, 16)
pad_h = 0
pad_w = 0
stride_h = 1
stride_w = 1
filter_h = 1
filter_w = 1
Data = tvm.placeholder(shape_in, name='FmapW', dtype="float16")
Weight = tvm.placeholder(shape_w, name='FilterW', dtype="float16")
bias_tensor = tvm.placeholder(
    (shape_w[1] * shape_w[2], ), name='Bias', dtype="float16")
res_tensor = dsl.conv(
    Data, Weight, {"bias_tensor": bias_tensor,
                   "pad_h": pad_h, "pad_w": pad_w,
                   "stride_h": stride_h, "stride_w": stride_w,
                   "filter_h": filter_h, "filter_w": filter_w,
                   "offset_a":0})