compute_five2four

此接口后续版本会废弃,新开发算子请不要使用此接口。

功能说明

把给定5-D “NC1HWC0”数据格式转换为4-D “NCHW”数据格式。

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

约束说明

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

支持的数据类型:float16。

函数原型

te.lang.cce.compute_five2four(input, raw_shape_4D)

参数说明

返回值

res_tensor:转换为4-D格式(N, C, H, W)后的tensor,tvm.tensor类型。

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

调用示例

import tvm
import te.lang.cce
raw_shape = (2, 32, 16, 128)
input_shape = (2,(32+15)//16,16,128,16)
in_dtype = "float16"
input = tvm.placeholder(input_shape, name='input', dtype=in_dtype)
res = te.lang.cce.compute_five2four(input, raw_shape)
# res.shape = (2, 32, 16, 128)