文档首页> Atlas 300应用(型号 3000)> TE API参考> compute接口> te.lang.cce.unsorted_segment_mean(tensor, segment_ids, num_segments, init_value=0)
更新时间:2021-03-18 GMT+08:00
分享

te.lang.cce.unsorted_segment_mean(tensor, segment_ids, num_segments, init_value=0)

使用数组segment_ids对tensor进行分段求均值。假设输入为data,输出为output,则output[i] = (1 / len(j...)) sum(data[j...]),其中“j...”是一个数组,“j...”中的元素j满足:segment_ids[j] == i。

如果某个下标i在segment_ids中没有出现,则output[i] = init_value。比如下图中,1在segment_ids中没有出现,则output[1] = 0。

如果segment_ids中某个值为负数,则对应位置的data中的值则丢弃。比如下图中,segment_ids[3] = -1,则data[3]的值被丢弃,不参与计算。

segment_ids其长度必须和data的第一维的长度相同。num_segments必须大于等于segment_ids的最大值加1。

支持的数据类型:float16、float32、int32。

该接口在segment_compute.py中定义。

参数说明

  • tensor:输入tensor,必须是如下数据类型之一:float16、float32、int32。
  • segment_ids:一维数组,对输入tensor进行分段的数组,其长度必须和输入tensor的第一维的长度相同。同时支持有序和无序。
  • num_segments:输出tensor的第一维的长度。其值必须大于等于segment_ids的最大值加1。
  • init_value:当segment_ids中某个下标不存在时,其输出的默认值。根据算子的实现来确定。默认值为0。

返回值:

res_tensor:表示计算后的tensor。

调用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
import tvm
import te.lang.cce 
shape = (5,1024) 
input_dtype = "float16"
data = tvm.placeholder(shape, name="data1", dtype=input_dtype) 
segment_ids = [1,1,5,5,5] 
num_segments = 6
res = te.lang.cce.unsorted_segment_mean(data, segment_ids, num_segments) 
# res.shape = (6,1024)
# res[0] = 0
# res[1] = (data[0] + data[1]) / 2
# res[2] = 0
# res[3] = 0
# res[4] = 0
# res[5] = (data[2] + data[3] + data[4]) / 3
分享:

    相关文档

    相关产品