<jupyter_start><jupyter_text>Save as checkpoint<jupyter_code>import tensorflow as tf
class Net(tf.keras.Model):
  def __init__(self):
    super(Net, self).__init__()
    self.l1 = tf.keras.layers.Dense(5)

  def call(self, x):
    return self.l1(x)
net = Net()
net.save_weights('easy_checkpoint')
%%bash
ls
def toy_dataset():
  inputs = tf.range(10.)[:, None] # None means shape=(10, 1)
  labels = inputs * 5. + tf.range(5.)[None, :]
  return tf.data.Dataset.from_tensor_slices(
    dict(x=inputs, y=labels)).repeat().batch(2)
def train_step(net, example, optimizer):
  with tf.GradientTape() as tape:
    output = net(example['x'])
    loss = tf.reduce_mean(tf.abs(output - example['y']))
  variables = net.trainable_variables
  gradients = tape.gradient(loss, variables)
  optimizer.apply_gradients(zip(gradients, variables))
  return loss
opt = tf.keras.optimizers.Adam(0.1)
dataset = toy_dataset()
iterator = iter(dataset)
ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=net, iterator=iterator)
manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)
# Save checkpoint at every step
def train_and_checkpoint(net, manager):
  ckpt.restore(manager.latest_checkpoint)
  if manager.latest_checkpoint:
    print('restored from {}'.format(manager.latest_checkpoint))
  else:
    print('from scratch')

  for _ in range(50):
    example = next(iterator)
    loss = train_step(net, example, opt)
    ckpt.step.assign_add(1) # number of steps
    if int(ckpt.step) % 10 == 0:
      save_path = manager.save()
      print('Saved ckpt for step {}: {}'.format(int(ckpt.step), save_path))
      print('loss {:1.2f}'.format(loss.numpy()))
train_and_checkpoint(net, manager)<jupyter_output><empty_output><jupyter_text>Restore and training<jupyter_code>opt = tf.keras.optimizers.Adam(0.1)
net = Net()
dataset = toy_dataset()
iterator = iter(dataset)
ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=net, iterator=iterator)
manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)

train_and_checkpoint(net, manager)
print(manager.checkpoints)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14896616541353383, 'fraction_numerical': 0.01268796992481203, 'mean_word_length': 5.298816568047338, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '624063', 'license': 'MIT', 'n_tokens_mistral': 781, 'n_tokens_neox': 759, 'n_words': 145, 'path': 'tensorflow/tf_model_save.ipynb', 'repo_name': 'satojkovic/runnable_tutorials', 'size': '2128'}