Translation Train

translation_train.py

import matplotlib.pyplot as plt
import numpy as np

def train(epochs, model, batch=64, shuffle=1000):
    loss_fcn = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True,
        reduction=tf.keras.losses.Reduction.NONE)
    opt = tf.keras.optimizers.Adam()
    losses = []
    ds = dataset.shuffle(shuffle).batch(batch).cache()
    for epoch in range(epochs):
        epoch_losses = []
        for eng_text, spa_text in ds:
            with tf.GradientTape() as tape:
                logits, expected, mask = model(eng_text, spa_text)
                loss = loss_fcn(expected, logits)
                loss = tf.ragged.boolean_mask(loss, mask)
                loss = tf.reduce_sum(loss) * (1. / batch)
                epoch_losses.append(loss.numpy())
                grads = tape.gradient(loss, model.trainable_weights)
                opt.apply_gradients(zip(grads, model.trainable_weights))
        losses.append(np.mean(epoch_losses))
        print('Trained epoch: {}; loss: {}'.format(epoch, losses[epoch]))
    plt.plot(losses)
    plt.xlabel('Epoch')
    plt.ylabel('Losses')