Training the model
There are a number of steps to be performed in training that require a custom training loop. First, let's define a method that executes one step of the training loop. This method is defined in the s2s-training.py file:
@tf.function
def train_step(inp, targ, enc_hidden, max_gradient_norm=5):
loss = 0
with tf.GradientTape() as tape:
# print("inside gradient tape")
enc_output, enc_hidden = encoder(inp, enc_hidden)
dec_hidden = enc_hidden
dec_input = tf.expand_dims([start] * BATCH_SIZE, 1)
# Teacher forcing - feeding the target as the next input
for t in range(1, targ.shape[1]):
# passing enc_output to the decoder
predictions, dec_hidden, _ = decoder(dec_input,
dec_hidden, enc_output)
loss += s2s.loss_function(targ[:, t], predictions)
# using teacher forcing...