How to save trained models?

I would like to save my trained models for future uses, especially as I am considering building my PhD predictors based on this library. Is there any way to add the .save attribute to the demos/models.py models to achieve something similar to what it is explained here?
Many thanks

[alvarosg]

Here’s some examples for check-pointing and model storage using Sonnet 2 and TF2: https://github.com/deepmind/sonnet#tensorflow-checkpointing

[alvarosg]

Because graph_nets is not built on top of Keras, saving and restoring models is slightly different. To save the model you can need to use a tf.train.Saver and saver.save, and to restore it, you should build the tensorflow graph in the same way, and then use saver.restore. See example below:

def get_input_graphs():
  # Some function that returns a graphs.GraphsTuple

def build_and_connect_model(input_graphs):
  graph_network = modules.GraphNetwork(
      edge_model_fn=lambda: snt.Linear(output_size=4),
      node_model_fn=lambda: snt.Linear(output_size=4),
      global_model_fn=lambda: snt.Linear(output_size=4))
  output_graphs = graph_network(input_graphs)
  return graph_network, output_graphs

def log_variables(sess, variables):
  vars_out = sess.run(variables)
  print([(var.name, var_out.flatten()[:][0]) 
         for var, var_out in zip(variables, vars_out)])

  
# Saving it.
tf.reset_default_graph()
input_graphs = get_input_graphs()
graph_net, output_graphs = build_and_connect_model(input_graphs)
initializer = tf.global_variables_initializer()

saver = tf.train.Saver()
with tf.Session() as sess:
  sess.run(initializer)
  saver.save(sess, "/tmp/model")
  log_variables(sess, graph_net.variables)

# Reloading it later.
tf.reset_default_graph()
input_graphs = get_input_graphs()
graph_net, output_graphs = build_and_connect_model(input_graphs)
saver = tf.train.Saver()
with tf.Session() as sess:
  saver.restore(sess, "/tmp/model")
  log_variables(sess, graph_net.variables)

Hope this helps!