Transformer model for language understanding

View on TensorFlow.org Run in Google Colab View source on GitHub Download notebook

This tutorial trains a Transformer model to translate Portuguese to English. This is an advanced example that assumes knowledge of text generation and attention.

The core idea behind the Transformer model is self-attention—the ability to attend to different positions of the input sequence to compute a representation of that sequence. Transformer creates stacks of self-attention layers and is explained below in the sections Scaled dot product attention and Multi-head attention.

A transformer model handles variable-sized input using stacks of self-attention layers instead of RNNs or CNNs. This general architecture has a number of advantages:

  • It make no assumptions about the temporal/spatial relationships across the data. This is ideal for processing a set of objects (for example, StarCraft units).
  • Layer outputs can be calculated in parallel, instead of a series like an RNN.
  • Distant items can affect each other's output without passing through many RNN-steps, or convolution layers (see Scene Memory Transformer for example).
  • It can learn long-range dependencies. This is a challenge in many sequence tasks.

The downsides of this architecture are:

  • For a time-series, the output for a time-step is calculated from the entire history instead of only the inputs and current hidden-state. This may be less efficient.
  • If the input does have a temporal/spatial relationship, like text, some positional encoding must be added or the model will effectively see a bag of words.

After training the model in this notebook, you will be able to input a Portuguese sentence and return the English translation.

Attention heatmap

import tensorflow_datasets as tfds
import tensorflow as tf

import time
import numpy as np
import matplotlib.pyplot as plt

Setup input pipeline

Use TFDS to load the Portugese-English translation dataset from the TED Talks Open Translation Project.

This dataset contains approximately 50000 training examples, 1100 validation examples, and 2000 test examples.

examples, metadata = tfds.load('ted_hrlr_translate/pt_to_en', with_info=True,
                               as_supervised=True)
train_examples, val_examples = examples['train'], examples['validation']
Downloading and preparing dataset ted_hrlr_translate/pt_to_en/1.0.0 (download: 124.94 MiB, generated: Unknown size, total: 124.94 MiB) to /home/kbuilder/tensorflow_datasets/ted_hrlr_translate/pt_to_en/1.0.0...

HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Dl Completed...', max=1.0, style=Progre…
HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Dl Size...', max=1.0, style=ProgressSty…
HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Extraction completed...', max=1.0, styl…







HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))
Shuffling and writing examples to /home/kbuilder/tensorflow_datasets/ted_hrlr_translate/pt_to_en/1.0.0.incompleteJ9C394/ted_hrlr_translate-train.tfrecord

HBox(children=(FloatProgress(value=0.0, max=51785.0), HTML(value='')))
HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))
Shuffling and writing examples to /home/kbuilder/tensorflow_datasets/ted_hrlr_translate/pt_to_en/1.0.0.incompleteJ9C394/ted_hrlr_translate-validation.tfrecord

HBox(children=(FloatProgress(value=0.0, max=1193.0), HTML(value='')))
HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))
Shuffling and writing examples to /home/kbuilder/tensorflow_datasets/ted_hrlr_translate/pt_to_en/1.0.0.incompleteJ9C394/ted_hrlr_translate-test.tfrecord

HBox(children=(FloatProgress(value=0.0, max=1803.0), HTML(value='')))
Dataset ted_hrlr_translate downloaded and prepared to /home/kbuilder/tensorflow_datasets/ted_hrlr_translate/pt_to_en/1.0.0. Subsequent calls will reuse this data.

Create a custom subwords tokenizer from the training dataset.

tokenizer_en = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    (en.numpy() for pt, en in train_examples), target_vocab_size=2**13)

tokenizer_pt = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    (pt.numpy() for pt, en in train_examples), target_vocab_size=2**13)
sample_string = 'Transformer is awesome.'

tokenized_string = tokenizer_en.encode(sample_string)
print ('Tokenized string is {}'.format(tokenized_string))

original_string = tokenizer_en.decode(tokenized_string)
print ('The original string: {}'.format(original_string))

assert original_string == sample_string
Tokenized string is [7915, 1248, 7946, 7194, 13, 2799, 7877]
The original string: Transformer is awesome.

The tokenizer encodes the string by breaking it into subwords if the word is not in its dictionary.

for ts in tokenized_string:
  print ('{} ----> {}'.format(ts, tokenizer_en.decode([ts])))
7915 ----> T
1248 ----> ran
7946 ----> s
7194 ----> former 
13 ----> is 
2799 ----> awesome
7877 ----> .

BUFFER_SIZE = 20000
BATCH_SIZE = 64

Add a start and end token to the input and target.

def encode(lang1, lang2):
  lang1 = [tokenizer_pt.vocab_size] + tokenizer_pt.encode(
      lang1.numpy()) + [tokenizer_pt.vocab_size+1]

  lang2 = [tokenizer_en.vocab_size] + tokenizer_en.encode(
      lang2.numpy()) + [tokenizer_en.vocab_size+1]
  
  return lang1, lang2

You want to use Dataset.map to apply this function to each element of the dataset. Dataset.map runs in graph mode.

  • Graph tensors do not have a value.
  • In graph mode you can only use TensorFlow Ops and functions.

So you can't .map this function directly: You need to wrap it in a tf.py_function. The tf.py_function will pass regular tensors (with a value and a .numpy() method to access it), to the wrapped python function.

def tf_encode(pt, en):
  result_pt, result_en = tf.py_function(encode, [pt, en], [tf.int64, tf.int64])
  result_pt.set_shape([None])
  result_en.set_shape([None])

  return result_pt, result_en
MAX_LENGTH = 40
def filter_max_length(x, y, max_length=MAX_LENGTH):
  return tf.logical_and(tf.size(x) <= max_length,
                        tf.size(y) <= max_length)
train_dataset = train_examples.map(tf_encode)
train_dataset = train_dataset.filter(filter_max_length)
# cache the dataset to memory to get a speedup while reading from it.
train_dataset = train_dataset.cache()
train_dataset = train_dataset.shuffle(BUFFER_SIZE).padded_batch(BATCH_SIZE)
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)


val_dataset = val_examples.map(tf_encode)
val_dataset = val_dataset.filter(filter_max_length).padded_batch(BATCH_SIZE)
pt_batch, en_batch = next(iter(val_dataset))
pt_batch, en_batch
(<tf.Tensor: shape=(64, 38), dtype=int64, numpy=
 array([[8214,  342, 3032, ...,    0,    0,    0],
        [8214,   95,  198, ...,    0,    0,    0],
        [8214, 4479, 7990, ...,    0,    0,    0],
        ...,
        [8214,  584,   12, ...,    0,    0,    0],
        [8214,   59, 1548, ...,    0,    0,    0],
        [8214,  118,   34, ...,    0,    0,    0]])>,
 <tf.Tensor: shape=(64, 40), dtype=int64, numpy=
 array([[8087,   98,   25, ...,    0,    0,    0],
        [8087,   12,   20, ...,    0,    0,    0],
        [8087,   12, 5453, ...,    0,    0,    0],
        ...,
        [8087,   18, 2059, ...,    0,    0,    0],
        [8087,   16, 1436, ...,    0,    0,    0],
        [8087,   15,   57, ...,    0,    0,    0]])>)

Positional encoding

Since this model doesn't contain any recurrence or convolution, positional encoding is added to give the model some information about the relative position of the words in the sentence.

The positional encoding vector is added to the embedding vector. Embeddings represent a token in a d-dimensional space where tokens with similar meaning will be closer to each other. But the embeddings do not encode the relative position of words in a sentence. So after adding the positional encoding, words will be closer to each other based on the similarity of their meaning and their position in the sentence, in the d-dimensional space.

See the notebook on positional encoding to learn more about it. The formula for calculating the positional encoding is as follows:

$$\Large{PE_{(pos, 2i)} = sin(pos / 10000^{2i / d_{model} })} $$
$$\Large{PE_{(pos, 2i+1)} = cos(pos / 10000^{2i / d_{model} })} $$
def get_angles(pos, i, d_model):
  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))
  return pos * angle_rates
def positional_encoding(position, d_model):
  angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)
  
  # apply sin to even indices in the array; 2i
  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
  
  # apply cos to odd indices in the array; 2i+1
  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    
  pos_encoding = angle_rads[np.newaxis, ...]
    
  return tf.cast(pos_encoding, dtype=tf.float32)
pos_encoding = positional_encoding(50, 512)
print (pos_encoding.shape)

plt.pcolormesh(pos_encoding[0], cmap='RdBu')
plt.xlabel('Depth')
plt.xlim((0, 512))
plt.ylabel('Position')
plt.colorbar()
plt.show()
(1, 50, 512)

png

Masking

Mask all the pad tokens in the batch of sequence. It ensures that the model does not treat padding as the input. The mask indicates where pad value 0 is present: it outputs a 1 at those locations, and a 0 otherwise.

def create_padding_mask(seq):
  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)
  
  # add extra dimensions to add the padding
  # to the attention logits.
  return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)
x = tf.constant([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]])
create_padding_mask(x)
<tf.Tensor: shape=(3, 1, 1, 5), dtype=float32, numpy=
array([[[[0., 0., 1., 1., 0.]]],


       [[[0., 0., 0., 1., 1.]]],


       [[[1., 1., 1., 0., 0.]]]], dtype=float32)>

The look-ahead mask is used to mask the future tokens in a sequence. In other words, the mask indicates which entries should not be used.

This means that to predict the third word, only the first and second word will be used. Similarly to predict the fourth word, only the first, second and the third word will be used and so on.

def create_look_ahead_mask(size):
  mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
  return mask  # (seq_len, seq_len)
x = tf.random.uniform((1, 3))
temp = create_look_ahead_mask(x.shape[1])
temp
<tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[0., 1., 1.],
       [0., 0., 1.],
       [0., 0., 0.]], dtype=float32)>

Scaled dot product attention

scaled_dot_product_attention

The attention function used by the transformer takes three inputs: Q (query), K (key), V (value). The equation used to calculate the attention weights is:

$$\Large{Attention(Q, K, V) = softmax_k(\frac{QK^T}{\sqrt{d_k} }) V} $$

The dot-product attention is scaled by a factor of square root of the depth. This is done because for large values of depth, the dot product grows large in magnitude pushing the softmax function where it has small gradients resulting in a very hard softmax.

For example, consider that Q and K have a mean of 0 and variance of 1. Their matrix multiplication will have a mean of 0 and variance of dk. Hence, square root of dk is used for scaling (and not any other number) because the matmul of Q and K should have a mean of 0 and variance of 1, and you get a gentler softmax.

The mask is multiplied with -1e9 (close to negative infinity). This is done because the mask is summed with the scaled matrix multiplication of Q and K and is applied immediately before a softmax. The goal is to zero out these cells, and large negative inputs to softmax are near zero in the output.

def scaled_dot_product_attention(q, k, v, mask):
  """Calculate the attention weights.
  q, k, v must have matching leading dimensions.
  k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
  The mask has different shapes depending on its type(padding or look ahead) 
  but it must be broadcastable for addition.
  
  Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable 
          to (..., seq_len_q, seq_len_k). Defaults to None.
    
  Returns:
    output, attention_weights
  """

  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
  
  # scale matmul_qk
  dk = tf.cast(tf.shape(k)[-1], tf.float32)
  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

  # add the mask to the scaled tensor.
  if mask is not None:
    scaled_attention_logits += (mask * -1e9)  

  # softmax is normalized on the last axis (seq_len_k) so that the scores
  # add up to 1.
  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

  return output, attention_weights

As the softmax normalization is done on K, its values decide the amount of importance given to Q.

The output represents the multiplication of the attention weights and the V (value) vector. This ensures that the words you want to focus on are kept as-is and the irrelevant words are flushed out.

def print_out(q, k, v):
  temp_out, temp_attn = scaled_dot_product_attention(
      q, k, v, None)
  print ('Attention weights are:')
  print (temp_attn)
  print ('Output is:')
  print (temp_out)
np.set_printoptions(suppress=True)

temp_k = tf.constant([[10,0,0],
                      [0,10,0],
                      [0,0,10],
                      [0,0,10]], dtype=tf.float32)  # (4, 3)

temp_v = tf.constant([[   1,0],
                      [  10,0],
                      [ 100,5],
                      [1000,6]], dtype=tf.float32)  # (4, 2)

# This `query` aligns with the second `key`,
# so the second `value` is returned.
temp_q = tf.constant([[0, 10, 0]], dtype=tf.float32)  # (1, 3)
print_out(temp_q, temp_k, temp_v)
Attention weights are:
tf.Tensor([[0. 1. 0. 0.]], shape=(1, 4), dtype=float32)
Output is:
tf.Tensor([[10.  0.]], shape=(1, 2), dtype=float32)

# This query aligns with a repeated key (third and fourth), 
# so all associated values get averaged.
temp_q = tf.constant([[0, 0, 10]], dtype=tf.float32)  # (1, 3)
print_out(temp_q, temp_k, temp_v)
Attention weights are:
tf.Tensor([[0.  0.  0.5 0.5]], shape=(1, 4), dtype=float32)
Output is:
tf.Tensor([[550.    5.5]], shape=(1, 2), dtype=float32)

# This query aligns equally with the first and second key, 
# so their values get averaged.
temp_q = tf.constant([[10, 10, 0]], dtype=tf.float32)  # (1, 3)
print_out(temp_q, temp_k, temp_v)
Attention weights are:
tf.Tensor([[0.5 0.5 0.  0. ]], shape=(1, 4), dtype=float32)
Output is:
tf.Tensor([[5.5 0. ]], shape=(1, 2), dtype=float32)

Pass all the queries together.

temp_q = tf.constant([[0, 0, 10], [0, 10, 0], [10, 10, 0]], dtype=tf.float32)  # (3, 3)
print_out(temp_q, temp_k, temp_v)
Attention weights are:
tf.Tensor(
[[0.  0.  0.5 0.5]
 [0.  1.  0.  0. ]
 [0.5 0.5 0.  0. ]], shape=(3, 4), dtype=float32)
Output is:
tf.Tensor(
[[550.    5.5]
 [ 10.    0. ]
 [  5.5   0. ]], shape=(3, 2), dtype=float32)

Multi-head attention

multi-head attention

Multi-head attention consists of four parts:

  • Linear layers and split into heads.
  • Scaled dot-product attention.
  • Concatenation of heads.
  • Final linear layer.

Each multi-head attention block gets three inputs; Q (query), K (key), V (value). These are put through linear (Dense) layers and split up into multiple heads.

The scaled_dot_product_attention defined above is applied to each head (broadcasted for efficiency). An appropriate mask must be used in the attention step. The attention output for each head is then concatenated (using tf.transpose, and tf.reshape) and put through a final Dense layer.

Instead of one single attention head, Q, K, and V are split into multiple heads because it allows the model to jointly attend to information at different positions from different representational spaces. After the split each head has a reduced dimensionality, so the total computation cost is the same as a single head attention with full dimensionality.

class MultiHeadAttention(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads):
    super(MultiHeadAttention, self).__init__()
    self.num_heads = num_heads
    self.d_model = d_model
    
    assert d_model % self.num_heads == 0
    
    self.depth = d_model // self.num_heads
    
    self.wq = tf.keras.layers.Dense(d_model)
    self.wk = tf.keras.layers.Dense(d_model)
    self.wv = tf.keras.layers.Dense(d_model)
    
    self.dense = tf.keras.layers.Dense(d_model)
        
  def split_heads(self, x, batch_size):
    """Split the last dimension into (num_heads, depth).
    Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
    """
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
    return tf.transpose(x, perm=[0, 2, 1, 3])
    
  def call(self, v, k, q, mask):
    batch_size = tf.shape(q)[0]
    
    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)
    
    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)
    
    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
    scaled_attention, attention_weights = scaled_dot_product_attention(
        q, k, v, mask)
    
    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention, 
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)
        
    return output, attention_weights

Create a MultiHeadAttention layer to try out. At each location in the sequence, y, the MultiHeadAttention runs all 8 attention heads across all other locations in the sequence, returning a new vector of the same length at each location.

temp_mha = MultiHeadAttention(d_model=512, num_heads=8)
y = tf.random.uniform((1, 60, 512))  # (batch_size, encoder_sequence, d_model)
out, attn = temp_mha(y, k=y, q=y, mask=None)
out.shape, attn.shape
(TensorShape([1, 60, 512]), TensorShape([1, 8, 60, 60]))

Point wise feed forward network

Point wise feed forward network consists of two fully-connected layers with a ReLU activation in between.

def point_wise_feed_forward_network(d_model, dff):
  return tf.keras.Sequential([
      tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)
      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)
  ])
sample_ffn = point_wise_feed_forward_network(512, 2048)
sample_ffn(tf.random.uniform((64, 50, 512))).shape
TensorShape([64, 50, 512])

Encoder and decoder

transformer

The transformer model follows the same general pattern as a standard sequence to sequence with attention model.

  • The input sentence is passed through N encoder layers that generates an output for each word/token in the sequence.
  • The decoder attends on the encoder's output and its own input (self-attention) to predict the next word.

Encoder layer

Each encoder layer consists of sublayers:

  1. Multi-head attention (with padding mask)
  2. Point wise feed forward networks.

Each of these sublayers has a residual connection around it followed by a layer normalization. Residual connections help in avoiding the vanishing gradient problem in deep networks.

The output of each sublayer is LayerNorm(x + Sublayer(x)). The normalization is done on the d_model (last) axis. There are N encoder layers in the transformer.

class EncoderLayer(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads, dff, rate=0.1):
    super(EncoderLayer, self).__init__()

    self.mha = MultiHeadAttention(d_model, num_heads)
    self.ffn = point_wise_feed_forward_network(d_model, dff)

    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    
    self.dropout1 = tf.keras.layers.Dropout(rate)
    self.dropout2 = tf.keras.layers.Dropout(rate)
    
  def call(self, x, training, mask):

    attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)
    attn_output = self.dropout1(attn_output, training=training)
    out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)
    
    ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)
    ffn_output = self.dropout2(ffn_output, training=training)
    out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)
    
    return out2
sample_encoder_layer = EncoderLayer(512, 8, 2048)

sample_encoder_layer_output = sample_encoder_layer(
    tf.random.uniform((64, 43, 512)), False, None)

sample_encoder_layer_output.shape  # (batch_size, input_seq_len, d_model)
TensorShape([64, 43, 512])

Decoder layer

Each decoder layer consists of sublayers:

  1. Masked multi-head attention (with look ahead mask and padding mask)
  2. Multi-head attention (with padding mask). V (value) and K (key) receive the encoder output as inputs. Q (query) receives the output from the masked multi-head attention sublayer.
  3. Point wise feed forward networks

Each of these sublayers has a residual connection around it followed by a layer normalization. The output of each sublayer is LayerNorm(x + Sublayer(x)). The normalization is done on the d_model (last) axis.

There are N decoder layers in the transformer.

As Q receives the output from decoder's first attention block, and K receives the encoder output, the attention weights represent the importance given to the decoder's input based on the encoder's output. In other words, the decoder predicts the next word by looking at the encoder output and self-attending to its own output. See the demonstration above in the scaled dot product attention section.

class DecoderLayer(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads, dff, rate=0.1):
    super(DecoderLayer, self).__init__()

    self.mha1 = MultiHeadAttention(d_model, num_heads)
    self.mha2 = MultiHeadAttention(d_model, num_heads)

    self.ffn = point_wise_feed_forward_network(d_model, dff)
 
    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    
    self.dropout1 = tf.keras.layers.Dropout(rate)
    self.dropout2 = tf.keras.layers.Dropout(rate)
    self.dropout3 = tf.keras.layers.Dropout(rate)
    
    
  def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):
    # enc_output.shape == (batch_size, input_seq_len, d_model)

    attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)
    attn1 = self.dropout1(attn1, training=training)
    out1 = self.layernorm1(attn1 + x)
    
    attn2, attn_weights_block2 = self.mha2(
        enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)
    attn2 = self.dropout2(attn2, training=training)
    out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)
    
    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)
    ffn_output = self.dropout3(ffn_output, training=training)
    out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)
    
    return out3, attn_weights_block1, attn_weights_block2
sample_decoder_layer = DecoderLayer(512, 8, 2048)

sample_decoder_layer_output, _, _ = sample_decoder_layer(
    tf.random.uniform((64, 50, 512)), sample_encoder_layer_output, 
    False, None, None)

sample_decoder_layer_output.shape  # (batch_size, target_seq_len, d_model)
TensorShape([64, 50, 512])

Encoder

The Encoder consists of:

  1. Input Embedding
  2. Positional Encoding
  3. N encoder layers

The input is put through an embedding which is summed with the positional encoding. The output of this summation is the input to the encoder layers. The output of the encoder is the input to the decoder.

class Encoder(tf.keras.layers.Layer):
  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
               maximum_position_encoding, rate=0.1):
    super(Encoder, self).__init__()

    self.d_model = d_model
    self.num_layers = num_layers
    
    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
    self.pos_encoding = positional_encoding(maximum_position_encoding, 
                                            self.d_model)
    
    
    self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) 
                       for _ in range(num_layers)]
  
    self.dropout = tf.keras.layers.Dropout(rate)
        
  def call(self, x, training, mask):

    seq_len = tf.shape(x)[1]
    
    # adding embedding and position encoding.
    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]

    x = self.dropout(x, training=training)
    
    for i in range(self.num_layers):
      x = self.enc_layers[i](x, training, mask)
    
    return x  # (batch_size, input_seq_len, d_model)
sample_encoder = Encoder(num_layers=2, d_model=512, num_heads=8, 
                         dff=2048, input_vocab_size=8500,
                         maximum_position_encoding=10000)
temp_input = tf.random.uniform((64, 62), dtype=tf.int64, minval=0, maxval=200)

sample_encoder_output = sample_encoder(temp_input, training=False, mask=None)

print (sample_encoder_output.shape)  # (batch_size, input_seq_len, d_model)
(64, 62, 512)

Decoder

The Decoder consists of:

  1. Output Embedding
  2. Positional Encoding
  3. N decoder layers

The target is put through an embedding which is summed with the positional encoding. The output of this summation is the input to the decoder layers. The output of the decoder is the input to the final linear layer.

class Decoder(tf.keras.layers.Layer):
  def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,
               maximum_position_encoding, rate=0.1):
    super(Decoder, self).__init__()

    self.d_model = d_model
    self.num_layers = num_layers
    
    self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)
    self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)
    
    self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) 
                       for _ in range(num_layers)]
    self.dropout = tf.keras.layers.Dropout(rate)
    
  def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):

    seq_len = tf.shape(x)[1]
    attention_weights = {}
    
    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]
    
    x = self.dropout(x, training=training)

    for i in range(self.num_layers):
      x, block1, block2 = self.dec_layers[i](x, enc_output, training,
                                             look_ahead_mask, padding_mask)
      
      attention_weights['decoder_layer{}_block1'.format(i+1)] = block1
      attention_weights['decoder_layer{}_block2'.format(i+1)] = block2
    
    # x.shape == (batch_size, target_seq_len, d_model)
    return x, attention_weights
sample_decoder = Decoder(num_layers=2, d_model=512, num_heads=8, 
                         dff=2048, target_vocab_size=8000,
                         maximum_position_encoding=5000)
temp_input = tf.random.uniform((64, 26), dtype=tf.int64, minval=0, maxval=200)

output, attn = sample_decoder(temp_input, 
                              enc_output=sample_encoder_output, 
                              training=False,
                              look_ahead_mask=None, 
                              padding_mask=None)

output.shape, attn['decoder_layer2_block2'].shape
(TensorShape([64, 26, 512]), TensorShape([64, 8, 26, 62]))

Create the Transformer

Transformer consists of the encoder, decoder and a final linear layer. The output of the decoder is the input to the linear layer and its output is returned.

class Transformer(tf.keras.Model):
  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, 
               target_vocab_size, pe_input, pe_target, rate=0.1):
    super(Transformer, self).__init__()

    self.encoder = Encoder(num_layers, d_model, num_heads, dff, 
                           input_vocab_size, pe_input, rate)

    self.decoder = Decoder(num_layers, d_model, num_heads, dff, 
                           target_vocab_size, pe_target, rate)

    self.final_layer = tf.keras.layers.Dense(target_vocab_size)
    
  def call(self, inp, tar, training, enc_padding_mask, 
           look_ahead_mask, dec_padding_mask):

    enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)
    
    # dec_output.shape == (batch_size, tar_seq_len, d_model)
    dec_output, attention_weights = self.decoder(
        tar, enc_output, training, look_ahead_mask, dec_padding_mask)
    
    final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)
    
    return final_output, attention_weights
sample_transformer = Transformer(
    num_layers=2, d_model=512, num_heads=8, dff=2048, 
    input_vocab_size=8500, target_vocab_size=8000, 
    pe_input=10000, pe_target=6000)

temp_input = tf.random.uniform((64, 38), dtype=tf.int64, minval=0, maxval=200)
temp_target = tf.random.uniform((64, 36), dtype=tf.int64, minval=0, maxval=200)

fn_out, _ = sample_transformer(temp_input, temp_target, training=False, 
                               enc_padding_mask=None, 
                               look_ahead_mask=None,
                               dec_padding_mask=None)

fn_out.shape  # (batch_size, tar_seq_len, target_vocab_size)
TensorShape([64, 36, 8000])

Set hyperparameters

To keep this example small and relatively fast, the values for num_layers, d_model, and dff have been reduced.

The values used in the base model of transformer were; num_layers=6, d_model = 512, dff = 2048. See the paper for all the other versions of the transformer.

num_layers = 4
d_model = 128
dff = 512
num_heads = 8

input_vocab_size = tokenizer_pt.vocab_size + 2
target_vocab_size = tokenizer_en.vocab_size + 2
dropout_rate = 0.1

Optimizer

Use the Adam optimizer with a custom learning rate scheduler according to the formula in the paper.

$$\Large{lrate = d_{model}^{-0.5} * min(step{\_}num^{-0.5}, step{\_}num * warmup{\_}steps^{-1.5})}$$
class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
  def __init__(self, d_model, warmup_steps=4000):
    super(CustomSchedule, self).__init__()
    
    self.d_model = d_model
    self.d_model = tf.cast(self.d_model, tf.float32)

    self.warmup_steps = warmup_steps
    
  def __call__(self, step):
    arg1 = tf.math.rsqrt(step)
    arg2 = step * (self.warmup_steps ** -1.5)
    
    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
learning_rate = CustomSchedule(d_model)

optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, 
                                     epsilon=1e-9)
temp_learning_rate_schedule = CustomSchedule(d_model)

plt.plot(temp_learning_rate_schedule(tf.range(40000, dtype=tf.float32)))
plt.ylabel("Learning Rate")
plt.xlabel("Train Step")
Text(0.5, 0, 'Train Step')

png

Loss and metrics

Since the target sequences are padded, it is important to apply a padding mask when calculating the loss.

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')
def loss_function(real, pred):
  mask = tf.math.logical_not(tf.math.equal(real, 0))
  loss_ = loss_object(real, pred)

  mask = tf.cast(mask, dtype=loss_.dtype)
  loss_ *= mask
  
  return tf.reduce_sum(loss_)/tf.reduce_sum(mask)
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
    name='train_accuracy')

Training and checkpointing

transformer = Transformer(num_layers, d_model, num_heads, dff,
                          input_vocab_size, target_vocab_size, 
                          pe_input=input_vocab_size, 
                          pe_target=target_vocab_size,
                          rate=dropout_rate)
def create_masks(inp, tar):
  # Encoder padding mask
  enc_padding_mask = create_padding_mask(inp)
  
  # Used in the 2nd attention block in the decoder.
  # This padding mask is used to mask the encoder outputs.
  dec_padding_mask = create_padding_mask(inp)
  
  # Used in the 1st attention block in the decoder.
  # It is used to pad and mask future tokens in the input received by 
  # the decoder.
  look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])
  dec_target_padding_mask = create_padding_mask(tar)
  combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)
  
  return enc_padding_mask, combined_mask, dec_padding_mask

Create the checkpoint path and the checkpoint manager. This will be used to save checkpoints every n epochs.

checkpoint_path = "./checkpoints/train"

ckpt = tf.train.Checkpoint(transformer=transformer,
                           optimizer=optimizer)

ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)

# if a checkpoint exists, restore the latest checkpoint.
if ckpt_manager.latest_checkpoint:
  ckpt.restore(ckpt_manager.latest_checkpoint)
  print ('Latest checkpoint restored!!')

The target is divided into tar_inp and tar_real. tar_inp is passed as an input to the decoder. tar_real is that same input shifted by 1: At each location in tar_input, tar_real contains the next token that should be predicted.

For example, sentence = "SOS A lion in the jungle is sleeping EOS"

tar_inp = "SOS A lion in the jungle is sleeping"

tar_real = "A lion in the jungle is sleeping EOS"

The transformer is an auto-regressive model: it makes predictions one part at a time, and uses its output so far to decide what to do next.

During training this example uses teacher-forcing (like in the text generation tutorial). Teacher forcing is passing the true output to the next time step regardless of what the model predicts at the current time step.

As the transformer predicts each word, self-attention allows it to look at the previous words in the input sequence to better predict the next word.

To prevent the model from peaking at the expected output the model uses a look-ahead mask.

EPOCHS = 20
# The @tf.function trace-compiles train_step into a TF graph for faster
# execution. The function specializes to the precise shape of the argument
# tensors. To avoid re-tracing due to the variable sequence lengths or variable
# batch sizes (the last batch is smaller), use input_signature to specify
# more generic shapes.

train_step_signature = [
    tf.TensorSpec(shape=(None, None), dtype=tf.int64),
    tf.TensorSpec(shape=(None, None), dtype=tf.int64),
]

@tf.function(input_signature=train_step_signature)
def train_step(inp, tar):
  tar_inp = tar[:, :-1]
  tar_real = tar[:, 1:]
  
  enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)
  
  with tf.GradientTape() as tape:
    predictions, _ = transformer(inp, tar_inp, 
                                 True, 
                                 enc_padding_mask, 
                                 combined_mask, 
                                 dec_padding_mask)
    loss = loss_function(tar_real, predictions)

  gradients = tape.gradient(loss, transformer.trainable_variables)    
  optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))
  
  train_loss(loss)
  train_accuracy(tar_real, predictions)

Portuguese is used as the input language and English is the target language.

for epoch in range(EPOCHS):
  start = time.time()
  
  train_loss.reset_states()
  train_accuracy.reset_states()
  
  # inp -> portuguese, tar -> english
  for (batch, (inp, tar)) in enumerate(train_dataset):
    train_step(inp, tar)
    
    if batch % 50 == 0:
      print ('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(
          epoch + 1, batch, train_loss.result(), train_accuracy.result()))
      
  if (epoch + 1) % 5 == 0:
    ckpt_save_path = ckpt_manager.save()
    print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,
                                                         ckpt_save_path))
    
  print ('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1, 
                                                train_loss.result(), 
                                                train_accuracy.result()))

  print ('Time taken for 1 epoch: {} secs\n'.format(time.time() - start))
Epoch 1 Batch 0 Loss 9.0395 Accuracy 0.0000
Epoch 1 Batch 50 Loss 8.9895 Accuracy 0.0003
Epoch 1 Batch 100 Loss 8.8906 Accuracy 0.0129
Epoch 1 Batch 150 Loss 8.7889 Accuracy 0.0175
Epoch 1 Batch 200 Loss 8.6619 Accuracy 0.0199
Epoch 1 Batch 250 Loss 8.5080 Accuracy 0.0219
Epoch 1 Batch 300 Loss 8.3304 Accuracy 0.0248
Epoch 1 Batch 350 Loss 8.1425 Accuracy 0.0295
Epoch 1 Batch 400 Loss 7.9585 Accuracy 0.0336
Epoch 1 Batch 450 Loss 7.7906 Accuracy 0.0367
Epoch 1 Batch 500 Loss 7.6424 Accuracy 0.0395
Epoch 1 Batch 550 Loss 7.5072 Accuracy 0.0426
Epoch 1 Batch 600 Loss 7.3788 Accuracy 0.0462
Epoch 1 Batch 650 Loss 7.2570 Accuracy 0.0497
Epoch 1 Batch 700 Loss 7.1419 Accuracy 0.0532
Epoch 1 Loss 7.1374 Accuracy 0.0533
Time taken for 1 epoch: 62.02073836326599 secs

Epoch 2 Batch 0 Loss 5.4534 Accuracy 0.1119
Epoch 2 Batch 50 Loss 5.5217 Accuracy 0.1016
Epoch 2 Batch 100 Loss 5.4412 Accuracy 0.1049
Epoch 2 Batch 150 Loss 5.3887 Accuracy 0.1060
Epoch 2 Batch 200 Loss 5.3362 Accuracy 0.1086
Epoch 2 Batch 250 Loss 5.2912 Accuracy 0.1105
Epoch 2 Batch 300 Loss 5.2550 Accuracy 0.1122
Epoch 2 Batch 350 Loss 5.2164 Accuracy 0.1140
Epoch 2 Batch 400 Loss 5.1831 Accuracy 0.1156
Epoch 2 Batch 450 Loss 5.1495 Accuracy 0.1173
Epoch 2 Batch 500 Loss 5.1197 Accuracy 0.1190
Epoch 2 Batch 550 Loss 5.0918 Accuracy 0.1206
Epoch 2 Batch 600 Loss 5.0644 Accuracy 0.1219
Epoch 2 Batch 650 Loss 5.0389 Accuracy 0.1232
Epoch 2 Batch 700 Loss 5.0161 Accuracy 0.1243
Epoch 2 Loss 5.0152 Accuracy 0.1244
Time taken for 1 epoch: 32.39113259315491 secs

Epoch 3 Batch 0 Loss 4.6219 Accuracy 0.1457
Epoch 3 Batch 50 Loss 4.6302 Accuracy 0.1426
Epoch 3 Batch 100 Loss 4.6107 Accuracy 0.1422
Epoch 3 Batch 150 Loss 4.6029 Accuracy 0.1429
Epoch 3 Batch 200 Loss 4.5940 Accuracy 0.1433
Epoch 3 Batch 250 Loss 4.5803 Accuracy 0.1437
Epoch 3 Batch 300 Loss 4.5632 Accuracy 0.1446
Epoch 3 Batch 350 Loss 4.5452 Accuracy 0.1459
Epoch 3 Batch 400 Loss 4.5284 Accuracy 0.1469
Epoch 3 Batch 450 Loss 4.5106 Accuracy 0.1481
Epoch 3 Batch 500 Loss 4.4982 Accuracy 0.1492
Epoch 3 Batch 550 Loss 4.4817 Accuracy 0.1499
Epoch 3 Batch 600 Loss 4.4651 Accuracy 0.1507
Epoch 3 Batch 650 Loss 4.4497 Accuracy 0.1514
Epoch 3 Batch 700 Loss 4.4335 Accuracy 0.1523
Epoch 3 Loss 4.4332 Accuracy 0.1523
Time taken for 1 epoch: 33.1234769821167 secs

Epoch 4 Batch 0 Loss 4.0327 Accuracy 0.1506
Epoch 4 Batch 50 Loss 4.0972 Accuracy 0.1681
Epoch 4 Batch 100 Loss 4.0787 Accuracy 0.1688
Epoch 4 Batch 150 Loss 4.0721 Accuracy 0.1705
Epoch 4 Batch 200 Loss 4.0559 Accuracy 0.1716
Epoch 4 Batch 250 Loss 4.0435 Accuracy 0.1726
Epoch 4 Batch 300 Loss 4.0304 Accuracy 0.1738
Epoch 4 Batch 350 Loss 4.0147 Accuracy 0.1749
Epoch 4 Batch 400 Loss 3.9958 Accuracy 0.1760
Epoch 4 Batch 450 Loss 3.9825 Accuracy 0.1769
Epoch 4 Batch 500 Loss 3.9678 Accuracy 0.1780
Epoch 4 Batch 550 Loss 3.9548 Accuracy 0.1792
Epoch 4 Batch 600 Loss 3.9405 Accuracy 0.1801
Epoch 4 Batch 650 Loss 3.9285 Accuracy 0.1809
Epoch 4 Batch 700 Loss 3.9162 Accuracy 0.1817
Epoch 4 Loss 3.9155 Accuracy 0.1817
Time taken for 1 epoch: 32.326401472091675 secs

Epoch 5 Batch 0 Loss 3.5525 Accuracy 0.1823
Epoch 5 Batch 50 Loss 3.5787 Accuracy 0.1961
Epoch 5 Batch 100 Loss 3.5621 Accuracy 0.1982
Epoch 5 Batch 150 Loss 3.5585 Accuracy 0.1995
Epoch 5 Batch 200 Loss 3.5541 Accuracy 0.2003
Epoch 5 Batch 250 Loss 3.5448 Accuracy 0.2005
Epoch 5 Batch 300 Loss 3.5352 Accuracy 0.2010
Epoch 5 Batch 350 Loss 3.5290 Accuracy 0.2021
Epoch 5 Batch 400 Loss 3.5219 Accuracy 0.2033
Epoch 5 Batch 450 Loss 3.5097 Accuracy 0.2041
Epoch 5 Batch 500 Loss 3.5024 Accuracy 0.2048
Epoch 5 Batch 550 Loss 3.4955 Accuracy 0.2055
Epoch 5 Batch 600 Loss 3.4861 Accuracy 0.2061
Epoch 5 Batch 650 Loss 3.4759 Accuracy 0.2068
Epoch 5 Batch 700 Loss 3.4672 Accuracy 0.2074
Saving checkpoint for epoch 5 at ./checkpoints/train/ckpt-1
Epoch 5 Loss 3.4662 Accuracy 0.2074
Time taken for 1 epoch: 32.44874715805054 secs

Epoch 6 Batch 0 Loss 3.0704 Accuracy 0.2368
Epoch 6 Batch 50 Loss 3.1327 Accuracy 0.2219
Epoch 6 Batch 100 Loss 3.1488 Accuracy 0.2218
Epoch 6 Batch 150 Loss 3.1410 Accuracy 0.2229
Epoch 6 Batch 200 Loss 3.1370 Accuracy 0.2234
Epoch 6 Batch 250 Loss 3.1268 Accuracy 0.2241
Epoch 6 Batch 300 Loss 3.1278 Accuracy 0.2234
Epoch 6 Batch 350 Loss 3.1217 Accuracy 0.2241
Epoch 6 Batch 400 Loss 3.1166 Accuracy 0.2246
Epoch 6 Batch 450 Loss 3.1102 Accuracy 0.2249
Epoch 6 Batch 500 Loss 3.1044 Accuracy 0.2255
Epoch 6 Batch 550 Loss 3.1001 Accuracy 0.2262
Epoch 6 Batch 600 Loss 3.0924 Accuracy 0.2267
Epoch 6 Batch 650 Loss 3.0845 Accuracy 0.2273
Epoch 6 Batch 700 Loss 3.0787 Accuracy 0.2276
Epoch 6 Loss 3.0783 Accuracy 0.2276
Time taken for 1 epoch: 32.548781394958496 secs

Epoch 7 Batch 0 Loss 2.8529 Accuracy 0.2175
Epoch 7 Batch 50 Loss 2.7503 Accuracy 0.2472
Epoch 7 Batch 100 Loss 2.7347 Accuracy 0.2463
Epoch 7 Batch 150 Loss 2.7352 Accuracy 0.2463
Epoch 7 Batch 200 Loss 2.7336 Accuracy 0.2462
Epoch 7 Batch 250 Loss 2.7276 Accuracy 0.2462
Epoch 7 Batch 300 Loss 2.7213 Accuracy 0.2457
Epoch 7 Batch 350 Loss 2.7158 Accuracy 0.2460
Epoch 7 Batch 400 Loss 2.7076 Accuracy 0.2468
Epoch 7 Batch 450 Loss 2.7014 Accuracy 0.2473
Epoch 7 Batch 500 Loss 2.6990 Accuracy 0.2471
Epoch 7 Batch 550 Loss 2.6953 Accuracy 0.2476
Epoch 7 Batch 600 Loss 2.6911 Accuracy 0.2478
Epoch 7 Batch 650 Loss 2.6878 Accuracy 0.2485
Epoch 7 Batch 700 Loss 2.6842 Accuracy 0.2487
Epoch 7 Loss 2.6836 Accuracy 0.2487
Time taken for 1 epoch: 32.43076515197754 secs

Epoch 8 Batch 0 Loss 2.4104 Accuracy 0.2812
Epoch 8 Batch 50 Loss 2.4010 Accuracy 0.2605
Epoch 8 Batch 100 Loss 2.3850 Accuracy 0.2622
Epoch 8 Batch 150 Loss 2.3775 Accuracy 0.2639
Epoch 8 Batch 200 Loss 2.3792 Accuracy 0.2640
Epoch 8 Batch 250 Loss 2.3862 Accuracy 0.2647
Epoch 8 Batch 300 Loss 2.3810 Accuracy 0.2652
Epoch 8 Batch 350 Loss 2.3754 Accuracy 0.2661
Epoch 8 Batch 400 Loss 2.3730 Accuracy 0.2660
Epoch 8 Batch 450 Loss 2.3716 Accuracy 0.2664
Epoch 8 Batch 500 Loss 2.3708 Accuracy 0.2665
Epoch 8 Batch 550 Loss 2.3705 Accuracy 0.2666
Epoch 8 Batch 600 Loss 2.3691 Accuracy 0.2671
Epoch 8 Batch 650 Loss 2.3672 Accuracy 0.2671
Epoch 8 Batch 700 Loss 2.3648 Accuracy 0.2673
Epoch 8 Loss 2.3640 Accuracy 0.2672
Time taken for 1 epoch: 32.15506458282471 secs

Epoch 9 Batch 0 Loss 2.1365 Accuracy 0.2878
Epoch 9 Batch 50 Loss 2.1157 Accuracy 0.2781
Epoch 9 Batch 100 Loss 2.1149 Accuracy 0.2810
Epoch 9 Batch 150 Loss 2.1266 Accuracy 0.2805
Epoch 9 Batch 200 Loss 2.1286 Accuracy 0.2802
Epoch 9 Batch 250 Loss 2.1264 Accuracy 0.2810
Epoch 9 Batch 300 Loss 2.1335 Accuracy 0.2809
Epoch 9 Batch 350 Loss 2.1330 Accuracy 0.2807
Epoch 9 Batch 400 Loss 2.1335 Accuracy 0.2806
Epoch 9 Batch 450 Loss 2.1329 Accuracy 0.2801
Epoch 9 Batch 500 Loss 2.1330 Accuracy 0.2801
Epoch 9 Batch 550 Loss 2.1306 Accuracy 0.2807
Epoch 9 Batch 600 Loss 2.1316 Accuracy 0.2806
Epoch 9 Batch 650 Loss 2.1322 Accuracy 0.2805
Epoch 9 Batch 700 Loss 2.1334 Accuracy 0.2803
Epoch 9 Loss 2.1337 Accuracy 0.2803
Time taken for 1 epoch: 32.37309765815735 secs

Epoch 10 Batch 0 Loss 1.9493 Accuracy 0.3042
Epoch 10 Batch 50 Loss 1.8978 Accuracy 0.2927
Epoch 10 Batch 100 Loss 1.9074 Accuracy 0.2959
Epoch 10 Batch 150 Loss 1.9176 Accuracy 0.2957
Epoch 10 Batch 200 Loss 1.9288 Accuracy 0.2938
Epoch 10 Batch 250 Loss 1.9325 Accuracy 0.2944
Epoch 10 Batch 300 Loss 1.9358 Accuracy 0.2941
Epoch 10 Batch 350 Loss 1.9358 Accuracy 0.2934
Epoch 10 Batch 400 Loss 1.9399 Accuracy 0.2930
Epoch 10 Batch 450 Loss 1.9402 Accuracy 0.2930
Epoch 10 Batch 500 Loss 1.9439 Accuracy 0.2926
Epoch 10 Batch 550 Loss 1.9477 Accuracy 0.2927
Epoch 10 Batch 600 Loss 1.9495 Accuracy 0.2923
Epoch 10 Batch 650 Loss 1.9517 Accuracy 0.2921
Epoch 10 Batch 700 Loss 1.9545 Accuracy 0.2920
Saving checkpoint for epoch 10 at ./checkpoints/train/ckpt-2
Epoch 10 Loss 1.9544 Accuracy 0.2920
Time taken for 1 epoch: 32.52263617515564 secs

Epoch 11 Batch 0 Loss 1.8385 Accuracy 0.2756
Epoch 11 Batch 50 Loss 1.7509 Accuracy 0.3050
Epoch 11 Batch 100 Loss 1.7687 Accuracy 0.3043
Epoch 11 Batch 150 Loss 1.7738 Accuracy 0.3034
Epoch 11 Batch 200 Loss 1.7830 Accuracy 0.3024
Epoch 11 Batch 250 Loss 1.7903 Accuracy 0.3010
Epoch 11 Batch 300 Loss 1.7922 Accuracy 0.3020
Epoch 11 Batch 350 Loss 1.7933 Accuracy 0.3015
Epoch 11 Batch 400 Loss 1.7955 Accuracy 0.3019
Epoch 11 Batch 450 Loss 1.7969 Accuracy 0.3010
Epoch 11 Batch 500 Loss 1.7992 Accuracy 0.3010
Epoch 11 Batch 550 Loss 1.8046 Accuracy 0.3010
Epoch 11 Batch 600 Loss 1.8079 Accuracy 0.3011
Epoch 11 Batch 650 Loss 1.8139 Accuracy 0.3008
Epoch 11 Batch 700 Loss 1.8150 Accuracy 0.3008
Epoch 11 Loss 1.8153 Accuracy 0.3009
Time taken for 1 epoch: 32.28882932662964 secs

Epoch 12 Batch 0 Loss 1.6943 Accuracy 0.3033
Epoch 12 Batch 50 Loss 1.6268 Accuracy 0.3104
Epoch 12 Batch 100 Loss 1.6283 Accuracy 0.3108
Epoch 12 Batch 150 Loss 1.6467 Accuracy 0.3089
Epoch 12 Batch 200 Loss 1.6571 Accuracy 0.3094
Epoch 12 Batch 250 Loss 1.6659 Accuracy 0.3083
Epoch 12 Batch 300 Loss 1.6717 Accuracy 0.3075
Epoch 12 Batch 350 Loss 1.6762 Accuracy 0.3077
Epoch 12 Batch 400 Loss 1.6764 Accuracy 0.3079
Epoch 12 Batch 450 Loss 1.6820 Accuracy 0.3076
Epoch 12 Batch 500 Loss 1.6865 Accuracy 0.3070
Epoch 12 Batch 550 Loss 1.6903 Accuracy 0.3073
Epoch 12 Batch 600 Loss 1.6960 Accuracy 0.3070
Epoch 12 Batch 650 Loss 1.6980 Accuracy 0.3070
Epoch 12 Batch 700 Loss 1.7000 Accuracy 0.3068
Epoch 12 Loss 1.7000 Accuracy 0.3068
Time taken for 1 epoch: 33.22452735900879 secs

Epoch 13 Batch 0 Loss 1.3527 Accuracy 0.2981
Epoch 13 Batch 50 Loss 1.5465 Accuracy 0.3148
Epoch 13 Batch 100 Loss 1.5470 Accuracy 0.3171
Epoch 13 Batch 150 Loss 1.5488 Accuracy 0.3171
Epoch 13 Batch 200 Loss 1.5571 Accuracy 0.3159
Epoch 13 Batch 250 Loss 1.5654 Accuracy 0.3156
Epoch 13 Batch 300 Loss 1.5678 Accuracy 0.3153
Epoch 13 Batch 350 Loss 1.5733 Accuracy 0.3155
Epoch 13 Batch 400 Loss 1.5755 Accuracy 0.3151
Epoch 13 Batch 450 Loss 1.5789 Accuracy 0.3158
Epoch 13 Batch 500 Loss 1.5821 Accuracy 0.3155
Epoch 13 Batch 550 Loss 1.5873 Accuracy 0.3150
Epoch 13 Batch 600 Loss 1.5911 Accuracy 0.3147
Epoch 13 Batch 650 Loss 1.5964 Accuracy 0.3141
Epoch 13 Batch 700 Loss 1.5992 Accuracy 0.3137
Epoch 13 Loss 1.5993 Accuracy 0.3137
Time taken for 1 epoch: 32.32143759727478 secs

Epoch 14 Batch 0 Loss 1.3310 Accuracy 0.3165
Epoch 14 Batch 50 Loss 1.4425 Accuracy 0.3231
Epoch 14 Batch 100 Loss 1.4505 Accuracy 0.3199
Epoch 14 Batch 150 Loss 1.4529 Accuracy 0.3195
Epoch 14 Batch 200 Loss 1.4661 Accuracy 0.3200
Epoch 14 Batch 250 Loss 1.4713 Accuracy 0.3203
Epoch 14 Batch 300 Loss 1.4762 Accuracy 0.3202
Epoch 14 Batch 350 Loss 1.4794 Accuracy 0.3200
Epoch 14 Batch 400 Loss 1.4857 Accuracy 0.3201
Epoch 14 Batch 450 Loss 1.4909 Accuracy 0.3195
Epoch 14 Batch 500 Loss 1.4960 Accuracy 0.3196
Epoch 14 Batch 550 Loss 1.5001 Accuracy 0.3195
Epoch 14 Batch 600 Loss 1.5051 Accuracy 0.3191
Epoch 14 Batch 650 Loss 1.5113 Accuracy 0.3191
Epoch 14 Batch 700 Loss 1.5157 Accuracy 0.3190
Epoch 14 Loss 1.5160 Accuracy 0.3189
Time taken for 1 epoch: 32.350428104400635 secs

Epoch 15 Batch 0 Loss 1.3344 Accuracy 0.3513
Epoch 15 Batch 50 Loss 1.3384 Accuracy 0.3312
Epoch 15 Batch 100 Loss 1.3699 Accuracy 0.3307
Epoch 15 Batch 150 Loss 1.3751 Accuracy 0.3291
Epoch 15 Batch 200 Loss 1.3862 Accuracy 0.3281
Epoch 15 Batch 250 Loss 1.3939 Accuracy 0.3282
Epoch 15 Batch 300 Loss 1.4009 Accuracy 0.3278
Epoch 15 Batch 350 Loss 1.4024 Accuracy 0.3273
Epoch 15 Batch 400 Loss 1.4090 Accuracy 0.3273
Epoch 15 Batch 450 Loss 1.4131 Accuracy 0.3268
Epoch 15 Batch 500 Loss 1.4181 Accuracy 0.3267
Epoch 15 Batch 550 Loss 1.4249 Accuracy 0.3261
Epoch 15 Batch 600 Loss 1.4304 Accuracy 0.3254
Epoch 15 Batch 650 Loss 1.4364 Accuracy 0.3252
Epoch 15 Batch 700 Loss 1.4416 Accuracy 0.3247
Saving checkpoint for epoch 15 at ./checkpoints/train/ckpt-3
Epoch 15 Loss 1.4413 Accuracy 0.3246
Time taken for 1 epoch: 32.81221079826355 secs

Epoch 16 Batch 0 Loss 1.3265 Accuracy 0.3153
Epoch 16 Batch 50 Loss 1.3004 Accuracy 0.3337
Epoch 16 Batch 100 Loss 1.2959 Accuracy 0.3349
Epoch 16 Batch 150 Loss 1.3075 Accuracy 0.3334
Epoch 16 Batch 200 Loss 1.3131 Accuracy 0.3328
Epoch 16 Batch 250 Loss 1.3183 Accuracy 0.3321
Epoch 16 Batch 300 Loss 1.3269 Accuracy 0.3317
Epoch 16 Batch 350 Loss 1.3325 Accuracy 0.3319
Epoch 16 Batch 400 Loss 1.3384 Accuracy 0.3315
Epoch 16 Batch 450 Loss 1.3436 Accuracy 0.3315
Epoch 16 Batch 500 Loss 1.3501 Accuracy 0.3314
Epoch 16 Batch 550 Loss 1.3548 Accuracy 0.3310
Epoch 16 Batch 600 Loss 1.3606 Accuracy 0.3304
Epoch 16 Batch 650 Loss 1.3651 Accuracy 0.3299
Epoch 16 Batch 700 Loss 1.3732 Accuracy 0.3294
Epoch 16 Loss 1.3735 Accuracy 0.3294
Time taken for 1 epoch: 32.512186765670776 secs

Epoch 17 Batch 0 Loss 1.0709 Accuracy 0.3410
Epoch 17 Batch 50 Loss 1.2269 Accuracy 0.3357
Epoch 17 Batch 100 Loss 1.2358 Accuracy 0.3381
Epoch 17 Batch 150 Loss 1.2516 Accuracy 0.3365
Epoch 17 Batch 200 Loss 1.2600 Accuracy 0.3347
Epoch 17 Batch 250 Loss 1.2667 Accuracy 0.3353
Epoch 17 Batch 300 Loss 1.2734 Accuracy 0.3354
Epoch 17 Batch 350 Loss 1.2804 Accuracy 0.3350
Epoch 17 Batch 400 Loss 1.2856 Accuracy 0.3346
Epoch 17 Batch 450 Loss 1.2902 Accuracy 0.3336
Epoch 17 Batch 500 Loss 1.2966 Accuracy 0.3339
Epoch 17 Batch 550 Loss 1.3002 Accuracy 0.3338
Epoch 17 Batch 600 Loss 1.3051 Accuracy 0.3339
Epoch 17 Batch 650 Loss 1.3115 Accuracy 0.3337
Epoch 17 Batch 700 Loss 1.3157 Accuracy 0.3334
Epoch 17 Loss 1.3158 Accuracy 0.3334
Time taken for 1 epoch: 32.974884271621704 secs

Epoch 18 Batch 0 Loss 1.2210 Accuracy 0.3707
Epoch 18 Batch 50 Loss 1.1708 Accuracy 0.3449
Epoch 18 Batch 100 Loss 1.1921 Accuracy 0.3413
Epoch 18 Batch 150 Loss 1.1981 Accuracy 0.3408
Epoch 18 Batch 200 Loss 1.2064 Accuracy 0.3401
Epoch 18 Batch 250 Loss 1.2156 Accuracy 0.3396
Epoch 18 Batch 300 Loss 1.2215 Accuracy 0.3398
Epoch 18 Batch 350 Loss 1.2261 Accuracy 0.3396
Epoch 18 Batch 400 Loss 1.2308 Accuracy 0.3397
Epoch 18 Batch 450 Loss 1.2341 Accuracy 0.3392
Epoch 18 Batch 500 Loss 1.2408 Accuracy 0.3386
Epoch 18 Batch 550 Loss 1.2474 Accuracy 0.3383
Epoch 18 Batch 600 Loss 1.2533 Accuracy 0.3379
Epoch 18 Batch 650 Loss 1.2576 Accuracy 0.3376
Epoch 18 Batch 700 Loss 1.2619 Accuracy 0.3376
Epoch 18 Loss 1.2619 Accuracy 0.3376
Time taken for 1 epoch: 33.297264099121094 secs

Epoch 19 Batch 0 Loss 1.1546 Accuracy 0.3243
Epoch 19 Batch 50 Loss 1.1284 Accuracy 0.3442
Epoch 19 Batch 100 Loss 1.1461 Accuracy 0.3441
Epoch 19 Batch 150 Loss 1.1553 Accuracy 0.3434
Epoch 19 Batch 200 Loss 1.1648 Accuracy 0.3449
Epoch 19 Batch 250 Loss 1.1692 Accuracy 0.3443
Epoch 19 Batch 300 Loss 1.1766 Accuracy 0.3432
Epoch 19 Batch 350 Loss 1.1808 Accuracy 0.3429
Epoch 19 Batch 400 Loss 1.1826 Accuracy 0.3423
Epoch 19 Batch 450 Loss 1.1860 Accuracy 0.3425
Epoch 19 Batch 500 Loss 1.1905 Accuracy 0.3422
Epoch 19 Batch 550 Loss 1.1962 Accuracy 0.3419
Epoch 19 Batch 600 Loss 1.2026 Accuracy 0.3414
Epoch 19 Batch 650 Loss 1.2078 Accuracy 0.3411
Epoch 19 Batch 700 Loss 1.2133 Accuracy 0.3408
Epoch 19 Loss 1.2138 Accuracy 0.3408
Time taken for 1 epoch: 32.29413819313049 secs

Epoch 20 Batch 0 Loss 1.0577 Accuracy 0.3539
Epoch 20 Batch 50 Loss 1.0679 Accuracy 0.3507
Epoch 20 Batch 100 Loss 1.0881 Accuracy 0.3484
Epoch 20 Batch 150 Loss 1.1035 Accuracy 0.3503
Epoch 20 Batch 200 Loss 1.1091 Accuracy 0.3508
Epoch 20 Batch 250 Loss 1.1151 Accuracy 0.3493
Epoch 20 Batch 300 Loss 1.1226 Accuracy 0.3488
Epoch 20 Batch 350 Loss 1.1280 Accuracy 0.3489
Epoch 20 Batch 400 Loss 1.1348 Accuracy 0.3475
Epoch 20 Batch 450 Loss 1.1390 Accuracy 0.3475
Epoch 20 Batch 500 Loss 1.1456 Accuracy 0.3465
Epoch 20 Batch 550 Loss 1.1500 Accuracy 0.3454
Epoch 20 Batch 600 Loss 1.1569 Accuracy 0.3453
Epoch 20 Batch 650 Loss 1.1625 Accuracy 0.3449
Epoch 20 Batch 700 Loss 1.1682 Accuracy 0.3444
Saving checkpoint for epoch 20 at ./checkpoints/train/ckpt-4
Epoch 20 Loss 1.1685 Accuracy 0.3444
Time taken for 1 epoch: 32.60363984107971 secs


Evaluate

The following steps are used for evaluation:

  • Encode the input sentence using the Portuguese tokenizer (tokenizer_pt). Moreover, add the start and end token so the input is equivalent to what the model is trained with. This is the encoder input.
  • The decoder input is the start token == tokenizer_en.vocab_size.
  • Calculate the padding masks and the look ahead masks.
  • The decoder then outputs the predictions by looking at the encoder output and its own output (self-attention).
  • Select the last word and calculate the argmax of that.
  • Concatentate the predicted word to the decoder input as pass it to the decoder.
  • In this approach, the decoder predicts the next word based on the previous words it predicted.
def evaluate(inp_sentence):
  start_token = [tokenizer_pt.vocab_size]
  end_token = [tokenizer_pt.vocab_size + 1]
  
  # inp sentence is portuguese, hence adding the start and end token
  inp_sentence = start_token + tokenizer_pt.encode(inp_sentence) + end_token
  encoder_input = tf.expand_dims(inp_sentence, 0)
  
  # as the target is english, the first word to the transformer should be the
  # english start token.
  decoder_input = [tokenizer_en.vocab_size]
  output = tf.expand_dims(decoder_input, 0)
    
  for i in range(MAX_LENGTH):
    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
        encoder_input, output)
  
    # predictions.shape == (batch_size, seq_len, vocab_size)
    predictions, attention_weights = transformer(encoder_input, 
                                                 output,
                                                 False,
                                                 enc_padding_mask,
                                                 combined_mask,
                                                 dec_padding_mask)
    
    # select the last word from the seq_len dimension
    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)

    predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
    
    # return the result if the predicted_id is equal to the end token
    if predicted_id == tokenizer_en.vocab_size+1:
      return tf.squeeze(output, axis=0), attention_weights
    
    # concatentate the predicted_id to the output which is given to the decoder
    # as its input.
    output = tf.concat([output, predicted_id], axis=-1)

  return tf.squeeze(output, axis=0), attention_weights
def plot_attention_weights(attention, sentence, result, layer):
  fig = plt.figure(figsize=(16, 8))
  
  sentence = tokenizer_pt.encode(sentence)
  
  attention = tf.squeeze(attention[layer], axis=0)
  
  for head in range(attention.shape[0]):
    ax = fig.add_subplot(2, 4, head+1)
    
    # plot the attention weights
    ax.matshow(attention[head][:-1, :], cmap='viridis')

    fontdict = {'fontsize': 10}
    
    ax.set_xticks(range(len(sentence)+2))
    ax.set_yticks(range(len(result)))
    
    ax.set_ylim(len(result)-1.5, -0.5)
        
    ax.set_xticklabels(
        ['<start>']+[tokenizer_pt.decode([i]) for i in sentence]+['<end>'], 
        fontdict=fontdict, rotation=90)
    
    ax.set_yticklabels([tokenizer_en.decode([i]) for i in result 
                        if i < tokenizer_en.vocab_size], 
                       fontdict=fontdict)
    
    ax.set_xlabel('Head {}'.format(head+1))
  
  plt.tight_layout()
  plt.show()
def translate(sentence, plot=''):
  result, attention_weights = evaluate(sentence)
  
  predicted_sentence = tokenizer_en.decode([i for i in result 
                                            if i < tokenizer_en.vocab_size])  

  print('Input: {}'.format(sentence))
  print('Predicted translation: {}'.format(predicted_sentence))
  
  if plot:
    plot_attention_weights(attention_weights, sentence, result, plot)
translate("este é um problema que temos que resolver.")
print ("Real translation: this is a problem we have to solve .")
Input: este é um problema que temos que resolver.
Predicted translation: this is a problem that we have to solve the whole problem of solving a problem .
Real translation: this is a problem we have to solve .

translate("os meus vizinhos ouviram sobre esta ideia.")
print ("Real translation: and my neighboring homes heard about this idea .")
Input: os meus vizinhos ouviram sobre esta ideia.
Predicted translation: my neighbors have heard about this idea .
Real translation: and my neighboring homes heard about this idea .

translate("vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram.")
print ("Real translation: so i 'll just share with you some stories very quickly of some magical things that have happened .")
Input: vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram.
Predicted translation: so i 'm going to share a little bit with you some stories of some magical things that had happened .
Real translation: so i 'll just share with you some stories very quickly of some magical things that have happened .

You can pass different layers and attention blocks of the decoder to the plot parameter.

translate("este é o primeiro livro que eu fiz.", plot='decoder_layer4_block2')
print ("Real translation: this is the first book i've ever done.")
Input: este é o primeiro livro que eu fiz.
Predicted translation: this is the first book i did .

png

Real translation: this is the first book i've ever done.

Summary

In this tutorial, you learned about positional encoding, multi-head attention, the importance of masking and how to create a transformer.

Try using a different dataset to train the transformer. You can also create the base transformer or transformer XL by changing the hyperparameters above. You can also use the layers defined here to create BERT and train state of the art models. Futhermore, you can implement beam search to get better predictions.