TensorNet是一个构建在TensorFlow之上针对广告推荐等大规模稀疏场景优化的分布式训练框架。TensorNet的目标是让所有使用TensorFlow的开发者可以快速的、方便的训练出稀疏参数超过百亿的超大模型。
在仅有CPU的集群中使用参数服务器的异步训练模式是训练模型速度最快的方法,TensorNet异步训练架构与TensorFlow的异步训练架构有很大的区别:
TensorNet最核心的优化是将模型的embedding tensor优化到了最小。
如下图所示,对于最简单的wide&deep模型,如果在一个广告系统中有3亿用户,那么就需要定义一个维度为3亿的embedding矩阵,在训练模型时需要在这个3亿维的矩阵上做embedding_lookup
得到当前batch内的用户的embedding信息,进而在embedding之上做更加复杂的操作。
显而易见,在高维稀疏场景下embedding矩阵太大,占用内存多。当特征较多的时候单机无法存储整个模型。
TensorNet使用一个较小的,可以容纳特征在一个batch内所有数据的embedding矩阵代替TensorFlow默认实现中需要定义的较大的embedding矩阵。
如下图所示,在batch_size设置为1024的场景下,对于用户id特征,在TensorNet中只需要定义一个维度为1024的embedding矩阵,TensorNet的主要处理步骤如下:
virtual sparse feature
。virtual sparse feature
作为模型的输入。从上述可见,TensorNet由于极大的减小了模型所需要的embedding矩阵,通过parameter server的方式使得稀疏特征的维度可以支持到接近无限维,从而可以极大的提升模型的刻画能力。
由于TensorNet只更改了模型的第一层,从而模型的inference也变得极其简单。
在使用TensorNet构造模型的时候,可以将模型切分为两部分,如下图所示,embedding_lookup_graph
只在离线训练时使用,在线inference时只需要将sparse embedding导出成字典供inference_graph
作为输入即可,具体的请参考:为inference准备——模型切分,使用XLA方式进行在线预估,sparse embedding字典导出系列文档。