Not able to restore attention model properly 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,
                                 encoder=encoder,
                                 decoder=decoder)

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


The training step is

EPOCHS = 50

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,
                                                   batch,
                                                   batch_loss.numpy()))
  # saving (checkpoint) the model every 3 epochs
  if (epoch + 1) % 3 == 0:
    manager.save()

  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

checkpoint.restore('ckpt-ckptnumber.index')


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)
you have to use entire things like ckpt.restore("./tf_ckpts/ckpt-10"). please check the https://www.tensorflow.org/guide/checkpoint.