[Question] How to properly visualize the model structure as well as save model?

I tried tf.keras.utils.plot_model

But it seems not to work… only print the model name, is it because we did not use Functional API here?

Also for saving the model, model.save() works but later the weights and summary can not be called even if I do loaded_model(inputs) first.

I am also a little bit confused of Model vs. Layer class. I see all tutorials here use Model class for both query and candidate tower, but -> https://www.tensorflow.org/guide/keras/custom_layers_and_models#the_model_class . It says if you do not need to call fit (which is the case for query and candidate tower), consider use layer class

[maciejkula]

There are two ways of saving model state in TensorFlow:

1. Using checkpoints for saving model variables during training.
2. Using tf.saved_model.save/load to store and load models for serving.

Here are the Keras docs for checkpointing (1). For (2), you can call tf.saved_model.save on your trained Model instance.

[jasonzyx]

Below is how I save & load the model/index

scann = tfrs.layers.factorized_top_k.ScaNN(
    model.query_model, 
    num_reordering_candidates=100
)

scann.index_from_dataset(
    products.batch(128).map(
        lambda x: (x['product_token'] + ': ' + x['product_name'], # idx identifier
                   model.product_model({                          # transform embeddings 
                       "product_name": x['product_name'],
                       "product_description": x["product_description"]
                   }))
    )
)

path = './your_model_name'
tf.saved_model.save(
    scann,
    path,
    options=tf.saved_model.SaveOptions(namespace_whitelist=["Scann"])
)

loaded = tf.saved_model.load(path)

_, titles = loaded({
    "query_text": np.array(["hello world"])
})
for title in titles[0]:
    print(title)