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)