Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Question] How to serve TF2 SOK model in Triton Inference and convert it to ONNX? #422

Closed
tuanavu opened this issue Sep 25, 2023 · 1 comment
Labels
question Further information is requested

Comments

@tuanavu
Copy link

tuanavu commented Sep 25, 2023

Details

I'm currently working with an existing TensorFlow 2 (TF2) model and the SparseOperationsKit. This set up allows me to utilize the SparseEmbedding Layer of the SOK toolkit. However, I've found that I have to define the sok_model and tf_model separately for training.

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():   
    sok_model = SOKModel(
        dense_feature_stats=dense_feature_stats, 
        trainable_sparse_feature_vocab_dict=trainable_sparse_feature_vocab_dict, 
        pretrained_sparse_feature_info_map=pretrained_sparse_feature_info_map,
        dense_dim=dense_dim,
        sparse_dim=sparse_dim,
    )
    tf_model = TFModel(
        all_feature_names=all_feature_names,
        pretrained_sparse_feature_info_map=pretrained_sparse_feature_info_map,
        dense_dim=dense_dim,
        sparse_dim=sparse_dim,
    )
    
    sok_opt = sok.optimizers.Adam()
    tf_opt = tf.keras.optimizers.Adam()

After training the new TF2 model with SOK, I found that I need to export both the sok_model and the tf_model separately.

tf_model.save(save_dir)
saver = sok.Saver()
saver.dump_to_file()

The resulting outputs are as follows:

  • sok_model: This results in a collection of files named EmbeddingVariable_*_keys.file and EmbeddingVariable_*_values.file.
  • tf2 model: This exports saved_model.pb, variables files.

When I need to execute a local test prediction request, I have to load both models independently. I then call the inference_step as follows:

# Load the model
sok_model.load_pretrained_embedding_table()

tf_model = tf.saved_model.load(save_dir)

# Inference steps
@tf.function(experimental_relax_shapes=True, reduce_retracing=True)
def inference_step(inputs):
    return tf_model(sok_model(inputs, training=False), training=False)

# Call inference
res = inference_step(inputs)

Questions

  • Model Serving: My goal is to deploy this model on the Triton Inference Server. I'm seeking guidance or examples that could streamline this process. I'm also curious about the ideal structure for this deployment - would treating it as an ensemble model that includes both sok and TensorFlow 2 backends be beneficial? In terms of backends, which would be the optimal choice - HugeCTR, TensorFlow 2, or another option? If there are any resources or guides that could assist me in this situation, I'd appreciate the pointers. For HugeCTR, it seems necessary to export the model graph; I'm wondering how I can accomplish the same with this TensorFlow 2 model that utilizes the SOK toolkit.
  • Model Conversion to ONNX: According to the Hierarchical Parameter Server Demo, HugeCTR can load both the sparse and dense models and convert them to a single ONNX model. I'm wondering how I can perform a similar conversion for this merlin-tensorflow model that uses the SOK toolkit and exports both the sparse and dense model.

Environment details

  • Merlin version: nvcr.io/nvidia/merlin/merlin-tensorflow:23.02
@tuanavu tuanavu added the question Further information is requested label Sep 25, 2023
@KingsleyLiu-NV
Copy link
Collaborator

Hi @tuanavu, there are two solutions of deploying models trained with TF2 SOK:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
question Further information is requested
Projects
None yet
Development

No branches or pull requests

2 participants