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, 2020 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, 2020 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, 2020 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, 2020 at 7:25
  • $\begingroup$ Yes. Thats my mistake. Check the checkpoints once. $\endgroup$
    – Uday
    Apr 13, 2020 at 7:26

1 Answer 1


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


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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