I am referring this article on building an attention model using tensorflow.

I am trying to train a similar model on my dataset using google colab. Due to the session limit of colab and my large dataset, I need to save the model state and restore it to resume training.

However, I am not able to restore the model upon saving the parameters. I have saved the input and target tokenizers, model checkpoint and even the input and output tensors. However, every time I use checkpoint.restore and resume training the model it resumes training with a high loss(equal to random weights).

I always test my model before saving using the translate function on some test data and it generates a one line summary. However, when I restore the model and run some sample data on the translate function, I only get a single tag as output (as if it is a newly initialised model).

Here is my code

checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(optimizer=optimizer,

manager = tf.train.CheckpointManager(checkpoint, 'checkpoint_dir', max_to_keep=1)

The training step is


for epoch in range(EPOCHS):
  start = time.time()

  enc_hidden = encoder.initialize_hidden_state()
  total_loss = 0

  for (batch, (inp, targ)) in tqdm(enumerate(dataset.take(steps_per_epoch))):
    batch_loss = train_step(inp, targ, enc_hidden)
    total_loss += batch_loss

    if batch % 100 == 0:
      print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,
  # saving (checkpoint) the model every 3 epochs
  if (epoch + 1) % 3 == 0:

  print('Epoch {} Loss {:.4f}'.format(epoch + 1,
                                      total_loss / steps_per_epoch))
  print('Time taken for 1 epoch {} sec\n'.format(time.time() - start))

I restore by doing


I save the tokenizers (both input and output) using pickle

with open('inp_tokenizer.pickle', 'wb') as handle:
      pickle.dump(inp_lang_tokenizer, handle, protocol=pickle.HIGHEST_PROTOCOL)

I save the tensors using numpy.save()

np.save('X.npy', input_tensor)
  • 1
    $\begingroup$ Post the sample code of how are you saving and loading the models. $\endgroup$ – Uday Apr 13 at 3:44
  • $\begingroup$ I have edited the question and added sample code. Most of the code is taken from the reference article itself. $\endgroup$ – Anush Kini Apr 13 at 6:55
  • $\begingroup$ execute this print(manager.checkpoints). You will get to know all the checkpoints that are there In the disk. I think there is no need for the .index/ I checked their train loop. They have written "enc_hidden = encoder.initialize_hidden_state()" so once after loading, if you execute that code, you will get initial weights again. $\endgroup$ – Uday Apr 13 at 7:09
  • $\begingroup$ But isn't enc_hidden = encoder.initialize_hidden_state() executed in every epoch?. This would mean that in every epoch the weights are reset to their initial weights again. $\endgroup$ – Anush Kini Apr 13 at 7:25
  • $\begingroup$ Yes. Thats my mistake. Check the checkpoints once. $\endgroup$ – Uday Apr 13 at 7:26

you have to use entire things like ckpt.restore("./tf_ckpts/ckpt-10"). please check the https://www.tensorflow.org/guide/checkpoint.

| improve this answer | |

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.