使用说明

接口定义

初始化EmbeddingLayerFWD。

EmbeddingLayerFWD(const int num_embeddings, const int embedding_dim, __fp16 *_weight,const int padding_idx = -1, const float max_norm = 1.0, const float norm_type = 2.0) ->void

表1 EmbeddingLayerFWD输入参数
参数名称	数据类型	描述	取值范围
num_embeddings	int	嵌入层中嵌入向量的数量。	[1, INT_MAX]
embedding_dim	int	每个嵌入向量的维度。	[1, INT_MAX]
_weight	_fp16*	权重。	推理阶段要求非空
padding_idx	int	指定填充标记的参数。	[INT_MIN, num_embeddings)
max_norm	float	指定权重向量的最大范数。	推理阶段不涉及
norm_type	float	指定权重向量的范数类型。	推理阶段不涉及

执行算子运算。

Run(int *indice, int dims) ->void

表2 Run函数输入参数
参数名称	数据类型	描述	取值范围
indice	int*	需要取出的向量索引组成的数组。	-
dims	int	目标向量维度数。	[1, INT_MAX]

验证EmbeddingLayerFWD的输入参数，并在算子构造过程中自动触发执行。

ValidateInput(const int num_embeddings, const int embedding_dim, __fp16 *_weight, const int padding_idx, const float max_norm, const float norm_type) ->KDNN::Status

表3 ValidateInput函数输入参数
参数名称	数据类型	描述	取值范围
num_embeddings	int	嵌入层中嵌入向量的数量。	[1, INT_MAX]
embedding_dim	int	每个嵌入向量的维度。	[1, INT_MAX]
_weight	_fp16*	权重。	推理阶段要求非空
padding_idx	int	指定填充标记的参数。	[INT_MIN, num_embeddings)
max_norm	float	指定权重向量的最大范数。	推理阶段不涉及
norm_type	float	指定权重向量的范数类型。	推理阶段不涉及

支持数据类型

Embedding支持fp16数据类型。（g_weight数组数据类型）_weight支持数据类型为fp16。

使用示例

将g_weight数组中对应索引数组indice中每个索引对应的向量取出，嵌入到固定大小的结果向量中。

__fp16 g_weight[10][3] = {
{0.1, 0.12, 0.15}, {0.21, 0.22, 0.25}, {0.11, 0.21, 0.32}, {0.77, 0.02, 0.05}, {0.54, 0.32, 0.08},
{0.14, 0.42, 0.25}, {0.31, 0.2, 0.07}, {0.22, 0.14, 0.03}, {0.17, 0.28, 0.32}, {0.44, 0.26, 0.23}
}; //使用fp16类型初始化向量数组
int num_embeddings = 10;
int embedding_dim = 3;
int num = 4;
// 构造算子
KDNN::EmbeddingLayerFWD embedding(num_embeddings, embedding_dim, (__fp16*)g_weight);
int indice[4] = {1, 2, 3, 7}; //索引数组
// 执行算子
__fp16 *out = embedding.Run(indice, num);

父主题： Embedding