Skip to content

Instantly share code, notes, and snippets.

@LysandreJik

LysandreJik/save_hub_checkpoint.py

Created November 8, 2019 16:19
Show Gist options
  • Star 2 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save LysandreJik/e85ee423a54fcc2d2d0664f04490eb0e to your computer and use it in GitHub Desktop.
Save LysandreJik/e85ee423a54fcc2d2d0664f04490eb0e to your computer and use it in GitHub Desktop.
Download ZIP
Save a HUB model to a checkpoint
Raw
save_hub_checkpoint.py
import tensorflow as tf
import tensorflow_hub as hub
model_size = 'xlarge'.upper()
version = 2
model = hub.Module("https://tfhub.dev/google/albert_{}/{}".format(model_size.lower(), version), trainable=False)
init = tf.global_variables_initializer()
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init)
path = "path_to_albert/albert"
save_path = saver.save(sess, path)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment