昇腾社区首页
中文
注册

aclnnEmbeddingBag

支持的产品型号

  • [object Object]Atlas 推理系列产品[object Object]。
  • [object Object]Atlas 训练系列产品[object Object]。
  • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]。

函数原型

每个算子分为undefined,必须先调用“aclnnEmbeddingBagGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用“aclnnEmbeddingBag”接口执行计算。

  • aclnnStatus aclnnEmbeddingBagGetWorkspaceSize(const aclTensor* weight, const aclTensor* indices,const aclTensor* offsets, bool scaleGradByFreq,int64_t mode, bool sparse, const aclTensor* perSampleWeights, bool includeLastOffset, int64_t paddingIdx, aclTensor* output, aclTensor* offset2bag, aclTensor* bagSize, aclTensor* maxIndices, uint64_t* workspaceSize, aclOpExecutor** executor)
  • aclnnStatus aclnnEmbeddingBag(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

功能说明

  • 算子功能:根据indices从weight中获得一组被聚合的数,然后根据offsets的偏移和mode指定的聚合模式对获取的数进行max、sum、mean聚合。其余参数则更细化了计算过程的控制。
  • shape推导方式如下: 假设:

    weight的shape为(numWeight, embeddingDim) indices的shape为(bagIndices) offsets的shape为(bagOffsets)

    • 当mode为sum模式:

      output的shape 为 includeLastOffset ? (bagOffsets - 1, embeddingDim) : (bagOffsets, embeddingDim) offset2bag的shape 为 (bagIndices,) bagSize的shape 为 includeLastOffset ? (bagOffsets - 1) : (bagOffsets,) maxIndices的shape 为 includeLastOffset ? (bagOffsets - 1) : (bagOffsets,)

    • 当mode为mean模式:

      output的shape 为 includeLastOffset? (bagOffsets - 1, embeddingDim) : (bagOffsets, embeddingDim) offset2bag的shape 为 (bagIndices,) bagSize的shape 为 includeLastOffset ? (bagOffsets - 1) : (bagOffsets,) maxIndices的shape 为 includeLastOffset ? (bagOffsets - 1) : (bagOffsets,)

    • 当mode为max模式:

      output的shape 为 includeLastOffset ? (bagOffsets - 1, embeddingDim) : (bagOffsets, embeddingDim) offset2bag的shape 为 (bagIndices,) bagSize的shape 为 includeLastOffset ? (bagOffsets - 1) : (bagOffsets,) maxIndices的shape 为 includeLastOffset ? (bagOffsets - 1, embeddingDim) : (bagOffsets, embeddingDim)

aclnnEmbeddingBagGetWorkspaceSize

  • 参数说明

    • weight(aclTensor*, 计算输入):词嵌入矩阵,包含所有词的嵌入向量,Device侧的aclTensor,shape支持2维,支持undefinedundefined支持ND。
      • [object Object]Atlas 训练系列产品[object Object]、[object Object]Atlas 推理系列产品[object Object]:数据类型支持FLOAT、FLOAT16。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]:数据类型支持FLOAT、FLOAT16、BFLOAT16。
    • indices(aclTensor*, 计算输入):包含索引的张量,指定要从weight中提取哪些词的嵌入向量,Device侧的aclTensor。数据类型支持UINT8、INT8、INT16、INT32、INT64,shape支持0-1维。支持undefinedundefined支持ND。
    • offsets(aclTensor*, 计算输入): 用于将 indices 分割成多个bag的偏移量张量,Device侧的aclTensor。数据类型支持UINT8、INT8、INT16、INT32、INT64。shape支持0-1维。
    • scaleGradByFreq(bool,计算输入): 用于控制是否根据词频缩放梯度,当scaleGradByFreq为true时,会根据词频对梯度进行缩放,当scaleGradByFreq为false时,则不会。
    • mode(int64_t, 计算输入):用于控制聚合模式,Host侧的整型。0表示sum聚合模式,1表示mean聚合模式,其他表示max聚合模式。
    • sparse(bool, 计算输入):用于控制稀疏模式,Host侧的bool类型。当为false时,表示weight非稀疏矩阵;当为true时,表示weight是稀疏矩阵。
    • perSampleWeights(aclTensor*, 计算输入): 指定样本权重,Device侧的aclTensor。shape支持1维,数据类型与weight一致,仅在sum模式下,可以不是nullptr,其他模式必须为nullptr。
      • [object Object]Atlas 训练系列产品[object Object]、[object Object]Atlas 推理系列产品[object Object]:数据类型支持FLOAT、FLOAT16。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]:数据类型支持FLOAT、FLOAT16、BFLOAT16。
    • includeLastOffset(bool, 计算输入):控制是否包含最后的偏移,Host侧的bool类型。当为false时,表示不包含最后的偏移;当为true时,表示包含最后的偏移。
    • paddingIdx(int64_t, 计算输入): 控制不参与计算的indices,Host侧的整型。
    • output(aclTensor*, 计算输出):词嵌入矩阵聚合后的结果,Device侧的aclTensor。数据类型与weight一致,shape支持2维。支持undefinedundefined支持ND。
      • [object Object]Atlas 训练系列产品[object Object]、[object Object]Atlas 推理系列产品[object Object]:数据类型支持FLOAT、FLOAT16。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]:数据类型支持FLOAT、FLOAT16、BFLOAT16。
    • offset2bag(aclTensor*, 计算输出):bag的起始偏移,Device侧的aclTensor。数据类型支持INT32、INT64,且offset2bag的数据类型和indices与offsets中精度高的一致,shape支持0-1维。支持undefinedundefined支持ND。
    • bagSize(aclTensor*, 计算输出): 每个bag的大小,Device侧的aclTensor。数据类型支持INT32、INT64,且offset2bag的数据类型和indices与offsets中精度高的一致,shape支持0-1维。支持undefinedundefined支持ND。
    • maxIndices(aclTensor*, 计算输出): 当mode为max时,词嵌入向量最大值所在的行,Device侧的aclTensor。数据类型支持INT32、INT64,且offset2bag的数据类型和indices与offsets中精度高的一致。当mode为max时,shape支持2维;当mode非max时,shape支持1维。支持undefinedundefined支持ND。
    • workspaceSize(uint64_t*, 出参):返回需要在Device侧申请的workspace大小。
    • executor(aclOpExecutor**, 出参):返回op执行器,包含了算子计算流程。
  • 返回值

    aclnnStatus:返回状态码,具体参见undefined

    [object Object]

aclnnEmbeddingBag

  • 参数说明

    • workspace(void*, 入参):在Device侧申请的workspace内存地址。
    • workspaceSize(uint64_t, 入参):在Device侧申请的workspace大小,由第一段接口aclnnEmbeddingBagGetWorkspaceSize获取。
    • executor(aclOpExecutor*, 入参):op执行器,包含了算子计算流程。
    • stream(aclrtStream, 入参):指定执行任务的AscendCL Stream流。
  • 返回值

    aclnnStatus:返回状态码,具体参见undefined

约束说明

sparse与scaleGradByFreq仅支持输入False。

调用示例

示例代码如下,仅供参考,具体编译和执行过程请参考undefined

[object Object]