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結果輸入模型.