功能描述

创建稀疏表。

函数原型

def create_table(key_dtype, dim, name, emb_initializer, device_vocabulary_size=1, host_vocabulary_size=0, ssd_vocabulary_size=0, ssd_data_path=(os.getcwd(),), is_save=True, is_dp=False, init_param=1., all2all_gradients_op=All2allGradientsOp.SUM_GRADIENTS.value, value_dtype=tf.float32, shard_num=1, fusion_optimizer_var=True, hashtable_threshold=0)

参数说明

参数名	类型	可选/必选	说明
key_dtype	TensorFlow的dtype类型	必选	稀疏特征键（key）数据类型。可选类型仅限于“tf.int64”和“tf.int32”。
dim	int tf.TensorShape	必选	嵌入层（embedding）维度。取值范围：[1, 8192]。如果dim需要设置为大于“512”的值，请保证内存和磁盘空间足够或者使用DDR模式，或者减小稀疏表的vocabulary size。如入参为tf.TensorShape类型，则要求其ndims为1，该值表示嵌入层维度。请根据服务器的实际配置进行设置。
name	str	必选	稀疏表表名，只能包含数字、字母和下划线。表名长度范围：[1, 100] 稀疏表表名需要保持唯一，不能重复。
emb_initializer	TensorFlow的初始化器类型	必选	嵌入层初始值生成器。
device_vocabulary_size	int	可选	Device侧嵌入层数量，默认值为“1”。取值范围：1~10亿。当设置超过25600000时，请保证内存和磁盘空间足够，或者开启片上内存显存侧动态扩容功能，或者减小稀疏表dim的大小。请根据服务器的实际配置进行设置。如果启用DDR/SSD存储，即host_vocabulary_size不为0，则需要device_vocabulary_size≥连续2个batch去重后的key个数，要求片上内存能存放至少2个batch数据，此时片上内存仅作为cache。
host_vocabulary_size	int	可选	Host侧DDR存储的嵌入层数量，默认值为“0”。取值范围为：0~10亿。取值为“0”时表示不开启Host侧DDR功能，如果不启用SSD存储，需要确保DDR能存储全量数据；不为“0”时表示开启，此时需要关闭片上内存显存侧的动态扩容，即“use_dynamic_expansion”=False，默认使用DDR内存侧的动态扩容模式。如果host_vocabulary_size设置为大于1亿的值，请保证内存和磁盘空间足够，或者减小稀疏表dim的大小。在动态扩容模式下（即use_dynamic_expansion=True时），默认使用片上内存作为唯一存储，此变量会被置“0”。超过单机内存时会出现OM。请根据服务器的实际配置进行设置。
ssd_vocabulary_size	int	可选	开启SSD存储embedding数据功能。默认值为“0”表示不开启。当值大于“0”时，要求host_vocabulary_size也大于“0”才能开启该功能。取值范围：0~10亿。在动态扩容模式下（即use_dynamic_expansion=True时），默认使用片上内存作为唯一存储，此变量会被置“0”。请根据服务器的实际配置进行设置。
ssd_data_path	list[str] Tuple[str]	可选	默认为当前运行脚本所在路径。当参数为空列表时，默认SSD存储路径为当前运行脚本所在路径。当列表非空且路径有效时，将按顺序存储到相应路径中。当路径对应磁盘空间不足时，会尝试下一个路径，直到所有磁盘空间不足时抛出异常。
is_save	bool	可选	是否保存embedding数据，默认值为“True”。取值范围： True：保存embedding数据 False：不保存embedding数据
is_dp	bool	可选	是否使能稀疏表数据并行功能，默认值为“False”。开启DP模式（is_dp=True）可以将表配置为数据并行，建议在存在10GB级小表，且当NPU侧为端到端瓶颈和稀疏表的NPU ALL2ALL通信量小于16MB时开启，能够带来一定的性能提升，稀疏通信收益15%左右。注意，混用DP和MP模式进行续训时，功能、精度不会受到影响，但不建议使用。
init_param	float	可选	Embedding初始化参数系数，默认值为“1.0”。取值范围：[-10, 10]。当init_param参数设置超过“1.0”或者小于“-1.0”时，建议减小batch_size，避免显存占用过大导致程序异常。
all2all_gradients_op	string	可选	分布式梯度回传后梯度聚合的方式，默认值为“sum_gradients”。 “sum_gradients”：所有rank的梯度相加。 “sum_gradients_and_div_by_ranksize”：所有rank梯度相加除以ranksize。
value_dtype	TensorFlow的dtype类型	可选	稀疏特征值（value）数据类型，默认值为“tf.float32”，仅支持“tf.float32”。
shard_num	int	可选	嵌入层分区数，默认值为“1”。取值范围：[1, 8192]。
fusion_optimizer_var	bool	可选	是否使用融合优化参数，默认值为“True”。取值范围： True：使用融合优化参数 False：不使用融合优化参数
hashtable_threshold	int	可选	哈希表阈值，高于阈值时使用哈希表，低于阈值时使用线性表，默认值为“0”。取值范围：[0, 2147483647]。

开启DDR/SSD模式必须确保自动改图模式也同时开启。
若使用片上内存显存侧的动态扩容方案，即init接口中传入use_dynamic_expansion=True，此时传入的host_vocabulary_size和ssd_vocabulary_size参数会被设置为“0”，即参数不生效，也可均不填。
若不使用片上内存显存侧的动态扩容方案，则根据host_vocabulary_size和ssd_vocabulary_size是否为0，来判断具体存储模式。
当“host_vocabulary_size”为“0”时，不开启Host侧DDR功能，不为“0”时开启。所有embedding表必须保持同时使用Host侧DDR功能或同时不使用Host侧DDR功能，即所有表“host_vocabulary_size”参数同时为“0”或同时不为“0”，否则进行参数校验时会报错，报错信息参考如下。
1
ValueError: The host-side DDR function of all tables must be used or not used at the same time. However, host voc size of each table is [].

返回值说明

成功：返回稀疏表实例。

通过返回的稀疏表实例可以访问该实例的两个方法说明如下：

方法

功能描述

函数原型

参数说明

返回值说明

使用示例

size方法

获取稀疏表的大小。

def size()

无

成功：返回稀疏表的大小。
失败：抛出异常

见create_table接口使用示例。

capacity方法

获取稀疏表的容量。

def capacity()

无

成功：返回稀疏表的容量。
失败：抛出异常

见create_table接口使用示例。

失败：抛出异常。

使用示例

import tensorflow as tf
from mx_rec.core.embedding import create_table
sparse_hashtable = create_table(key_dtype=tf.int32,
                                dim=tf.TensorShape([128]),
                                name="sparse_embeddings_table",
                                emb_initializer=tf.truncated_normal_initializer(),
                                device_vocabulary_size=24_000_000 * 8,
                                host_vocabulary_size=0)
table_size = sparse_hashtable.size()      # 获取返回稀疏表的使用大小
table_capacity = sparse_hashtable.capacity()   # 获取返回稀疏表的容量大小

参考资源

接口调用流程及示例，参见模型迁移与训练。