concat
功能说明
在指定轴上对输入的多个Tensor进行重新连接。
输入raw_tensors为多个Tensor,数据类型相同。
如果raw_tensors[i].shape = [D0, D1, ... Daxis(i), ...Dn],沿着轴axis连接后的结果的shape为:[D0, D1, ... Raxis, ...Dn]。
其中:Raxis = sum(Daxis(i))。
例如:
t1 = [[1, 2, 3], [4, 5, 6]] t2 = [[7, 8, 9], [10, 11, 12]] concat([t1, t2], 0) # [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] concat([t1, t2], 1) # [[1, 2, 3, 7, 8, 9], [4, 5, 6, 10, 11, 12]] # tensor t1的shape为 [2, 3] # tensor t2的shape为 [2, 3] concat([t1, t2], 0).shape # [4, 3] concat([t1, t2], 1).shape # [2, 6]
参数axis也可以为负数,表示从维度的最后开始计算,表示第axis + len(shape)跟轴。
例如:
t1 = [[[1, 2], [2, 3]], [[4, 4], [5, 3]]] t2 = [[[7, 4], [8, 4]], [[2, 10], [15, 11]]] concat([t1, t2], -1)
结果为:
[[[ 1, 2, 7, 4], [ 2, 3, 8, 4]], [[ 4, 4, 2, 10], [ 5, 3, 15, 11]]]
函数原型
concat(raw_tensors, axis)
参数说明
返回值
res_tensor:重新连接后的tensor,tvm.tensor类型。
约束说明
此接口暂不支持与其他TBE DSL计算接口混合使用。
支持的型号
Atlas 200/300/500 推理产品
Atlas 训练系列产品
Atlas 推理系列产品(Ascend 310P处理器)
Atlas 200/500 A2推理产品
Atlas A2训练系列产品/Atlas 800I A2推理产品
调用示例
from tbe import tvm from tbe import dsl shape1 = (64,128) shape2 = (64,128) input_dtype = "float16" data1 = tvm.placeholder(shape1, name="data1", dtype=input_dtype) data2 = tvm.placeholder(shape2, name="data1", dtype=input_dtype) data = [data1, data2] res = dsl.concat(data, 0) # res.shape = (128,128)
父主题: Tensor操作接口