文档首页> Atlas 300应用（型号 3000）> TE API参考> compute接口> te.lang.cce.conv(*args)
更新时间：2021-03-18 GMT+08:00
查看PDF
分享

te.lang.cce.conv(*args)

在给定4-D输入和filter的情况下计算2-D卷积。

要求输入Tensor和filter Tensor的格式均为NCHW。

接口可以支持bias，支持的数据类型float16。

该接口在conv_compute.py中定义。

参数说明

*args：是一个list，参数数量可变。
  • bias场景参数列表为（其中hasBias需要为True）：

    A, W, B, res_dtype, padh, padw, strideh, stridew, hasBias

  • 非bias场景参数列表为（其中hasBias需要为False）：

    A, W, res_dtype, padh, padw, strideh, stridew, hasBias

    其中：
    • A：输入Tensor，即卷积计算的feature map。
    • W：filter Tensor，即卷积计算的卷积核。
    • B：bias Tensor，即卷积计算的偏置。
    • res_dtype：输出Tensor的数据类型，即卷积计算结果的数据类型。
    • padh：padding的高，即卷积计算在feature map的H方向上的填充数。
    • padw：padding的宽，即卷积计算在feature map的W方向上的填充数。
    • strideh：H方向的步长，即卷积计算filter在feature map的H方向上移动的步长。
    • stridew：W方向的步长，即卷积计算filter在feature map的W方向上移动的步长。
    • hasBias：是否带bias。
  • 约束：假设feature map的shape为（Fn, Fc, Fh, Fw），filter的shape为（Wn, Wc, Wh, Ww），卷积结果的输出shape为（On, Oc, Oh, Ow），padh记为Ph，padw记为Pw，strideh记为Sh，stridew记为Sw，上述各参数之间需满足如下关系：
    • Fc = Wc
    • On = Fn
    • Oc = Wn
    • Oh = ((Fh + 2Ph - Wh) / Sh) + 1
    • Ow = ((Fw + 2Pw - Ww) / Sw) + 1

返回值

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

调用示例

import tvm
import te.lang.cce 
shape_in = (64,64) 
shape_w = (3,3) 
in_dtype = "float16"
A = tvm.placeholder(shape_in, name='A', dtype=in_dtype) 
W = tvm.placeholder(shape_w, name='W', dtype=in_dtype)  
b_shape = (shape_w[0], ) 
B = tvm.placeholder(b_shape, name='B', dtype=res_dtype)  
padh = 0
padw = 0  
strideh = 1
stridew = 1 
res = te.lang.cce.conv(A, W, B, in_dtype, padh, padw, strideh, stridew, True) 
# res = A * W + B
父主题： compute接口
分享：

    相关文档

    相关产品