tensorflow share embedding feature column

需求

多个feature column列需要进行embedding, 且embedding参数共享, 并共同更新embedding参数.

操作

tensorflow 1.13 中的embedding feature column, share embedding columns 仅支持从存储的tf 模型 ckpt文件中读取参数数组. 因此需要将不同数据源数据转化并存储到ckpt模型中. 然后指定路径+variable名进行加载.

# 存储模型
def gather_graph_variable(self):
    with tf.Session() as sess:
        for name, arr in self.tf_variable_dict.items():
            print("save {} embedding.".format(name))
            _ = tf.get_variable(name, initializer=arr)
        sess.run(tf.global_variables_initializer())
        saver = tf.train.Saver()
        save_path = saver.save(sess=sess, save_path=self.output_item_embedding_path)
        print("Model saved in path: %s" % save_path)
    

# share embedding
def get_shared_initialized_embedding(key, feature_name_list, vocabulary_path, ckpt_path, tensor_name, dtype, emb_dim=None,
                                     train_able=True):
    categorical_column_list = list()
    for feature_name in feature_name_list:
        feature_name = "{}_{}".format(key, feature_name)
        category_col = get_categorical_column_with_vocabulary_file(feature_name, vocabulary_path, dtype=dtype)
        categorical_column_list.append(category_col)
    columns = shared_embedding_columns(categorical_column_list, dimension=emb_dim, ckpt_to_load_from=ckpt_path,
                                       tensor_name_in_ckpt=tensor_name, trainable=train_able)
    return columns
    

ps

在开发过程中需要用到官方提供的 DNNLinearClassifier estimator. 因此特征转化操作全部需要使用预定义的feature_column进行(如果可以也可以开发自定义的feature column), 而不能在特征抽取过程中直接对input tensor进行操作.
如果采用自定义的estimator 则可以通过get_vairable, look_up 操作在特征处理中直接操作, 将tensor结果输入模型.