API Call Overview

Initialize the framework. Call the init API in main.py to input the parameters required for initializing the framework. For details about the related parameters, see init.

# nbatch function needs to be used together with the prefetch and host_vocabulary_size != 0
init(max_steps=max_steps,
     train_steps=TRAIN_steps,
     eval_steps=EVAL_STEPS,
     save_steps=SAVE_STEPS,
     use_dynamic=use_dynamic,
     use_dynamic_expansion=use_dynamic_expansion)

Define a dataset. Call the get_asc_insert_func API in main.py to create a dataset and preprocess the dataset. For details about the related parameters, see Parameters.

    if not MODIFY_GRAPH_FLAG:
        insert_fn = get_asc_insert_func(tgt_key_specs=feature_spec_list, is_training=is_training, dump_graph=dump_graph)
        dataset = dataset.map(insert_fn)
    dataset = dataset.prefetch(100)
    iterator = dataset.make_initializable_iterator()
    batch = iterator.get_next()
    return batch, iterator

Define an optimizer. Define an optimizer in optimizer.py. For details about supported optimizer types and related parameters, see Optimizers.

# coding: UTF-8
import logging
import tensorflow as tf
from mx_rec.optimizers.lazy_adam import create_hash_optimizer
from mx_rec.optimizers.lazy_adam_by_addr import create_hash_optimizer_by_address
from mx_rec.util.initialize import get_use_dynamic_expansion
def get_dense_and_sparse_optimizer(cfg):
    dense_optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=cfg.learning_rate)
    use_dynamic_expansion = get_use_dynamic_expansion()
    if use_dynamic_expansion:
        sparse_optimizer = create_hash_optimizer_by_address(learning_rate=cfg.learning_rate)
        logging.info("optimizer lazy_adam_by_addr")
    else:
        sparse_optimizer = create_hash_optimizer(learning_rate=cfg.learning_rate)
        logging.info("optimizer lazy_adam")
    return dense_optimizer, sparse_optimizer

Create a sparse table. Call the create_table API in main.py to create a sparse table and a sparse network layer. For details about the related parameters, see Parameters.

user_hashtable = create_table(key_dtype=tf.int64,
                              dim=tf.Tensorshape([cfg.user_hashtable_dim]),
                              name='user_table',
                              emb_initializer=tf.compat.v1.truncated_normal_initializer(),
                              device_vocabulary_size=cfg.user_vocab_size * 10,
                              host_vocabulary_size=0)  # cfg.user_vocab_size * 100, # for h2d test
item_hashtable = create_table(key_dtype=tf.int64,
                              dim=tf.Tensorshape([cfg.item_hashtable_dim]),
                              name='item_table',
                              emb_initializer=tf.compat.v1.truncated_normal_initializer(),
                              device_vocabulary_size=cfg.item_vocab_size * 10,
                              host_vocabulary_size=0)  # cfg.user_vocab_size * 100, # for h2d test

Create a computational graph. Import the sparse network layer and feature list, create a model computational graph, and call the sparse_lookup API in the computational graph to query features and calculate errors. For details about the related parameters, see Parameters.

def model_forward(input_list, batch, is_train, modify_graph, config_dict=None):
    embedding_list = []
    feature_list, hash_table_list, send_count_list = input_list
    for feature, hash_table, send_count in zip(feature_list, hash_table_list, send_count_list):
        access_and_evict_config = None
        if isinstance(config_dict, dict):
            access_and_evict_config = config_dict.get(hash_table.table_name)
        embedding = sparse_lookup(hash_table, feature, send_count, is_train=is_train,
                                  access_and_evict_config=access_and_evict_config,
                                  name=hash_table.table_name + "_lookup", modify_graph=modify_graph, batch=batch)
        reduced_embedding = tf.reduce_sum(embedding, axis=1, keepdims=False)
        embedding_list.append(reduced_embedding)
    my_model = MyModel()
    my_model(embedding_list, batch["label_0"], batch["label_1"])
    return my_model

Define the gradient calculation and optimization processes. Call the get_dense_and_sparse_variable API in main.py to obtain the parameters of the dense network layer and sparse network layer. Use the optimizer to calculate gradients and perform optimization. For details about the API, see get_dense_and_sparse_variable.

train_iterator, train_model = build_graph([user_hashtable, item_hashtable], is_train=True,
                                              feature_spec_list=train_feature_spec_list,
                                              config_dict=ACCESS_AND_EVICT, batch_number=cfg.batch_number)
eval_iterator, eval_model = build_graph([user_hashtable, item_hashtable], is_train=False,
                                            feature_spec_list=eval_feature_spec_list,
                                            config_dict=ACCESS_AND_EVICT, batch_number=cfg.batch_number)
dense_variables, sparse_variables = get_dense_and_sparse_variable()

Load and preprocess data. Call the modify_graph_and_start_emb_cache (graph modification mode) or start_asc_pipeline (non-graph modification mode) API in main.py to start the data pipeline. (In the sample code, if is used to determine whether to use MODIFY_GRAPH_FLAG to set the graph modification mode in the configuration file.) For details about the API, see modify_graph_and_start_emb_cache.

saver = tf.compat.v1.train.Saver()
if MODIFY_GRAPH_FLAG:
    logging.info("start to modifying graph")
    modify_graph_and_start_emb_cache(dump_graph=True)
else:
    start_asc_pipeline()
with tf.compat.v1.Session(config=sess_config(dump_data=False)) as sess:
    if MODIFY_GRAPH_FLAG:
        sess.run(get_initializer(True))
    else:
        sess.run(train_iterator.initializer)
    sess.run(tf.compat.v1.global_variables_initializer())
    EPOCH = 0
    if os.path.exists(f"./saved-model/sparse-model-{rank_id}-%d" % 0):
        saver.restore(sess, f"./saved-model/model-{rank_id}-%d" % 0)
    else:
        saver.save(sess, f"./saved-model/model-{rank_id}", global_step=0)
    for i in range(1, 201):
        logging.info(f"################    training at step {i}    ################")
        try:
            sess.run([train_ops, train_model.loss_list])
        except tf.errors.OutOfRangeError:
            logging.info(f"Encounter the end of Sequence for training.")
            break
        else:
            if i % TRAIN_INTERVAL == 0:
                EPOCH += 1
                evaluate()
            if i % SAVING_INTERVAL == 0:
                saver.save(sess, f"./saved-model/model-{rank_id}", global_step=i)
    saver.save(sess, f"./saved-model/model-{rank_id}", global_step=i)

Start session calculation and save the model during training by calling the saver API in main.py.

saver = tf.compat.v1.train.Saver()
if MODIFY_GRAPH_FLAG:
    logging.info("start to modifying graph")
    modify_graph_and_start_emb_cache(dump_graph=True)
else:
    start_asc_pipeline()
with tf.compat.v1.Session(config=sess_config(dump_data=False)) as sess:
    if MODIFY_GRAPH_FLAG:
        sess.run(get_initializer(True))
    else:
        sess.run(train_iterator.initializer)
    sess.run(tf.compat.v1.global_variables_initializer())
    EPOCH = 0
    if os.path.exists(f"./saved-model/sparse-model-{rank_id}-%d" % 0):
        saver.restore(sess, f"./saved-model/model-{rank_id}-%d" % 0)
    else:
        saver.save(sess, f"./saved-model/model-{rank_id}", global_step=0)
    for i in range(1, 201):
        logging.info(f"################    training at step {i}    ################")
        try:
            sess.run([train_ops, train_model.loss_list])
        except tf.errors.OutOfRangeError:
            logging.info(f"Encounter the end of Sequence for training.")
            break
        else:
            if i % TRAIN_INTERVAL == 0:
                EPOCH += 1
                evaluate()
            if i % SAVING_INTERVAL == 0:
                saver.save(sess, f"./saved-model/model-{rank_id}", global_step=i)
    saver.save(sess, f"./saved-model/model-{rank_id}", global_step=i)

Close the data flow to destroy resources. Call the terminate_config_initializer API in main.py to close the data flow and destroy resources. For details about the API, see terminate_config_initializer.
```
terminate_config_initializer()
logging.info("Demo done!")
```

Parent topic: Quick Start