Cette page a été traduite par l'API Cloud Translation.
Switch to English

Génération de texte avec un RNN

Voir sur TensorFlow.org Exécuter dans Google Colab Afficher la source sur GitHub Télécharger le cahier

Ce didacticiel montre comment générer du texte à l'aide d'un RNN basé sur des caractères. Nous travaillerons avec un ensemble de données de l'écriture de Shakespeare tirée de L'efficacité déraisonnable des réseaux de neurones récurrents d'Andrej Karpathy. Étant donné une séquence de caractères à partir de ces données ("Shakespear"), entraînez un modèle pour prédire le caractère suivant dans la séquence ("e"). Des séquences de texte plus longues peuvent être générées en appelant le modèle à plusieurs reprises.

Ce tutoriel comprend du code exécutable implémenté à l'aide de tf.keras et une exécution rapide . Voici un exemple de sortie lorsque le modèle de ce didacticiel s'est entraîné pendant 30 époques et a commencé par la chaîne «Q»:

QUEENE:
I had thought thou hadst a Roman; for the oracle,
Thus by All bids the man against the word,
Which are so weak of care, by old care done;
Your children were in your holy love,
And the precipitation through the bleeding throne.

BISHOP OF ELY:
Marry, and will, my lord, to weep in such a one were prettiest;
Yet now I was adopted heir
Of the world's lamentable day,
To watch the next way with his father with his face?

ESCALUS:
The cause why then we are all resolved more sons.

VOLUMNIA:
O, no, no, no, no, no, no, no, no, no, no, no, no, no, no, no, no, no, no, no, no, it is no sin it should be dead,
And love and pale as any will to that word.

QUEEN ELIZABETH:
But how long have I heard the soul for this world,
And show his hands of life be proved to stand.

PETRUCHIO:
I say he look'd on, if I must be content
To stay him from the fatal of our country's bliss.
His lordship pluck'd from this sentence then for prey,
And then let us twain, being the moon,
were she such a case as fills m

Bien que certaines phrases soient grammaticales, la plupart n'ont pas de sens. Le modèle n'a pas appris la signification des mots, mais considérez:

  • Le modèle est basé sur des caractères. Lorsque la formation a commencé, le modèle ne savait pas comment épeler un mot anglais, ou que les mots étaient même une unité de texte.

  • La structure de la sortie ressemble à une lecture - les blocs de texte commencent généralement par un nom de locuteur, en toutes lettres majuscules similaires à l'ensemble de données.

  • Comme démontré ci-dessous, le modèle est formé sur de petits lots de texte (100 caractères chacun), et est toujours capable de générer une séquence de texte plus longue avec une structure cohérente.

Installer

Importer TensorFlow et d'autres bibliothèques

 import tensorflow as tf

import numpy as np
import os
import time
 

Téléchargez le jeu de données Shakespeare

Modifiez la ligne suivante pour exécuter ce code sur vos propres données.

 path_to_file = tf.keras.utils.get_file('shakespeare.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt')
 
Downloading data from https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt
1122304/1115394 [==============================] - 0s 0us/step

Lisez les données

Tout d'abord, regardez dans le texte:

 # Read, then decode for py2 compat.
text = open(path_to_file, 'rb').read().decode(encoding='utf-8')
# length of text is the number of characters in it
print ('Length of text: {} characters'.format(len(text)))
 
Length of text: 1115394 characters

 # Take a look at the first 250 characters in text
print(text[:250])
 
First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

First Citizen:
You are all resolved rather to die than to famish?

All:
Resolved. resolved.

First Citizen:
First, you know Caius Marcius is chief enemy to the people.


 # The unique characters in the file
vocab = sorted(set(text))
print ('{} unique characters'.format(len(vocab)))
 
65 unique characters

Traiter le texte

Vectoriser le texte

Avant l'entraînement, nous devons mapper des chaînes à une représentation numérique. Créez deux tables de recherche: l'une mappant les caractères aux nombres et l'autre les nombres aux caractères.

 # Creating a mapping from unique characters to indices
char2idx = {u:i for i, u in enumerate(vocab)}
idx2char = np.array(vocab)

text_as_int = np.array([char2idx[c] for c in text])
 

Nous avons maintenant une représentation entière pour chaque caractère. Notez que nous avons mappé le caractère sous forme d'index de 0 à len(unique) .

 print('{')
for char,_ in zip(char2idx, range(20)):
    print('  {:4s}: {:3d},'.format(repr(char), char2idx[char]))
print('  ...\n}')
 
{
  '\n':   0,
  ' ' :   1,
  '!' :   2,
  '$' :   3,
  '&' :   4,
  "'" :   5,
  ',' :   6,
  '-' :   7,
  '.' :   8,
  '3' :   9,
  ':' :  10,
  ';' :  11,
  '?' :  12,
  'A' :  13,
  'B' :  14,
  'C' :  15,
  'D' :  16,
  'E' :  17,
  'F' :  18,
  'G' :  19,
  ...
}

 # Show how the first 13 characters from the text are mapped to integers
print ('{} ---- characters mapped to int ---- > {}'.format(repr(text[:13]), text_as_int[:13]))
 
'First Citizen' ---- characters mapped to int ---- > [18 47 56 57 58  1 15 47 58 47 64 43 52]

La tâche de prédiction

Étant donné un caractère ou une séquence de caractères, quel est le prochain caractère le plus probable? C'est la tâche que nous formons le modèle à effectuer. L'entrée du modèle sera une séquence de caractères, et nous entraînons le modèle pour prédire la sortie - le caractère suivant à chaque pas de temps.

Puisque les RNN conservent un état interne qui dépend des éléments vus précédemment, étant donné tous les caractères calculés jusqu'à ce moment, quel est le caractère suivant?

Créer des exemples et des objectifs de formation

Ensuite, divisez le texte en exemples de séquences. Chaque séquence d'entrée contiendra des caractères seq_length du texte.

Pour chaque séquence d'entrée, les cibles correspondantes contiennent la même longueur de texte, sauf décalées d'un caractère vers la droite.

Alors seq_length+1 le texte en morceaux de seq_length+1 . Par exemple, disons que seq_length vaut 4 et que notre texte est "Hello". La séquence d'entrée serait "Hell" et la séquence cible "ello".

Pour ce faire, utilisez d'abord la fonction tf.data.Dataset.from_tensor_slices pour convertir le vecteur de texte en un flux d'indices de caractères.

 # The maximum length sentence we want for a single input in characters
seq_length = 100
examples_per_epoch = len(text)//(seq_length+1)

# Create training examples / targets
char_dataset = tf.data.Dataset.from_tensor_slices(text_as_int)

for i in char_dataset.take(5):
  print(idx2char[i.numpy()])
 
F
i
r
s
t

La méthode par batch nous permet de convertir facilement ces caractères individuels en séquences de la taille souhaitée.

 sequences = char_dataset.batch(seq_length+1, drop_remainder=True)

for item in sequences.take(5):
  print(repr(''.join(idx2char[item.numpy()])))
 
'First Citizen:\nBefore we proceed any further, hear me speak.\n\nAll:\nSpeak, speak.\n\nFirst Citizen:\nYou '
'are all resolved rather to die than to famish?\n\nAll:\nResolved. resolved.\n\nFirst Citizen:\nFirst, you k'
"now Caius Marcius is chief enemy to the people.\n\nAll:\nWe know't, we know't.\n\nFirst Citizen:\nLet us ki"
"ll him, and we'll have corn at our own price.\nIs't a verdict?\n\nAll:\nNo more talking on't; let it be d"
'one: away, away!\n\nSecond Citizen:\nOne word, good citizens.\n\nFirst Citizen:\nWe are accounted poor citi'

Pour chaque séquence, dupliquez-la et déplacez-la pour former le texte d'entrée et cible en utilisant la méthode map pour appliquer une fonction simple à chaque lot:

 def split_input_target(chunk):
    input_text = chunk[:-1]
    target_text = chunk[1:]
    return input_text, target_text

dataset = sequences.map(split_input_target)
 

Imprimez les premiers exemples d'entrée et de valeurs cibles:

 for input_example, target_example in  dataset.take(1):
  print ('Input data: ', repr(''.join(idx2char[input_example.numpy()])))
  print ('Target data:', repr(''.join(idx2char[target_example.numpy()])))
 
Input data:  'First Citizen:\nBefore we proceed any further, hear me speak.\n\nAll:\nSpeak, speak.\n\nFirst Citizen:\nYou'
Target data: 'irst Citizen:\nBefore we proceed any further, hear me speak.\n\nAll:\nSpeak, speak.\n\nFirst Citizen:\nYou '

Chaque index de ces vecteurs est traité comme un pas de temps. Pour l'entrée au pas de temps 0, le modèle reçoit l'indice pour "F" et essaie de prédire l'indice pour "i" comme caractère suivant. Au pas de temps suivant, il fait la même chose mais le RNN considère le contexte de l'étape précédente en plus du caractère d'entrée actuel.

 for i, (input_idx, target_idx) in enumerate(zip(input_example[:5], target_example[:5])):
    print("Step {:4d}".format(i))
    print("  input: {} ({:s})".format(input_idx, repr(idx2char[input_idx])))
    print("  expected output: {} ({:s})".format(target_idx, repr(idx2char[target_idx])))
 
Step    0
  input: 18 ('F')
  expected output: 47 ('i')
Step    1
  input: 47 ('i')
  expected output: 56 ('r')
Step    2
  input: 56 ('r')
  expected output: 57 ('s')
Step    3
  input: 57 ('s')
  expected output: 58 ('t')
Step    4
  input: 58 ('t')
  expected output: 1 (' ')

Créer des lots d'entraînement

Nous avons utilisé tf.data pour diviser le texte en séquences gérables. Mais avant d'introduire ces données dans le modèle, nous devons mélanger les données et les regrouper en lots.

 # Batch size
BATCH_SIZE = 64

# Buffer size to shuffle the dataset
# (TF data is designed to work with possibly infinite sequences,
# so it doesn't attempt to shuffle the entire sequence in memory. Instead,
# it maintains a buffer in which it shuffles elements).
BUFFER_SIZE = 10000

dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)

dataset
 
<BatchDataset shapes: ((64, 100), (64, 100)), types: (tf.int64, tf.int64)>

Construire le modèle

Utilisez tf.keras.Sequential pour définir le modèle. Pour cet exemple simple, trois couches sont utilisées pour définir notre modèle:

  • tf.keras.layers.Embedding : La couche d'entrée. Une table de recherche entraînable qui mappera les nombres de chaque caractère à un vecteur avec des dimensions embedding_dim ;
  • tf.keras.layers.GRU : Un type de RNN avec des units=rnn_units taille units=rnn_units (Vous pouvez également utiliser une couche LSTM ici.)
  • tf.keras.layers.Dense : La couche de sortie, avec les sorties vocab_size .
 # Length of the vocabulary in chars
vocab_size = len(vocab)

# The embedding dimension
embedding_dim = 256

# Number of RNN units
rnn_units = 1024
 
 def build_model(vocab_size, embedding_dim, rnn_units, batch_size):
  model = tf.keras.Sequential([
    tf.keras.layers.Embedding(vocab_size, embedding_dim,
                              batch_input_shape=[batch_size, None]),
    tf.keras.layers.GRU(rnn_units,
                        return_sequences=True,
                        stateful=True,
                        recurrent_initializer='glorot_uniform'),
    tf.keras.layers.Dense(vocab_size)
  ])
  return model
 
 model = build_model(
  vocab_size = len(vocab),
  embedding_dim=embedding_dim,
  rnn_units=rnn_units,
  batch_size=BATCH_SIZE)
 

Pour chaque caractère, le modèle recherche l'incorporation, exécute le GRU un pas de temps avec l'incorporation comme entrée et applique la couche dense pour générer des logits prédisant la vraisemblance logarithmique du caractère suivant:

Un dessin des données traversant le modèle

Veuillez noter que nous choisissons ici le modèle séquentiel Keras car toutes les couches du modèle n'ont qu'une seule entrée et produisent une seule sortie. Si vous souhaitez récupérer et réutiliser les états de la couche RNN avec état, vous pouvez créer votre modèle avec l'API fonctionnelle Keras ou la sous-classification de modèle. Veuillez consulter le guide Keras RNN pour plus de détails.

Essayez le modèle

Exécutez maintenant le modèle pour voir qu'il se comporte comme prévu.

Vérifiez d'abord la forme de la sortie:

 for input_example_batch, target_example_batch in dataset.take(1):
  example_batch_predictions = model(input_example_batch)
  print(example_batch_predictions.shape, "# (batch_size, sequence_length, vocab_size)")
 
(64, 100, 65) # (batch_size, sequence_length, vocab_size)

Dans l'exemple ci-dessus, la longueur de séquence de l'entrée est de 100 mais le modèle peut être exécuté sur des entrées de n'importe quelle longueur:

 model.summary()
 
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding (Embedding)        (64, None, 256)           16640     
_________________________________________________________________
gru (GRU)                    (64, None, 1024)          3938304   
_________________________________________________________________
dense (Dense)                (64, None, 65)            66625     
=================================================================
Total params: 4,021,569
Trainable params: 4,021,569
Non-trainable params: 0
_________________________________________________________________

Pour obtenir des prédictions réelles du modèle, nous devons échantillonner à partir de la distribution de sortie, pour obtenir les indices de caractères réels. Cette distribution est définie par les logits sur le vocabulaire des caractères.

Essayez-le pour le premier exemple du lot:

 sampled_indices = tf.random.categorical(example_batch_predictions[0], num_samples=1)
sampled_indices = tf.squeeze(sampled_indices,axis=-1).numpy()
 

Cela nous donne, à chaque pas de temps, une prédiction du prochain index de caractère:

 sampled_indices
 
array([36, 17, 45, 51, 54, 26, 37, 30, 12,  5, 30, 33, 22, 17, 46, 17,  9,
       59, 14, 11, 48, 38, 39, 62, 37,  8, 19, 23, 16, 11, 16, 35, 11, 54,
        9, 31,  5, 42, 38, 24, 44, 30, 60, 56, 31, 45, 56, 44, 47, 64, 54,
       26, 15, 28, 32, 44, 38, 60, 39, 15, 29, 37, 41, 40, 14,  4, 27, 56,
       14, 54, 45,  7, 14, 28, 25,  5, 34,  1, 15, 16, 19,  6, 63, 31, 11,
        6, 23,  3, 47, 33, 59, 54, 20,  4, 51, 12, 35, 20,  1, 37])

Décodez-les pour voir le texte prédit par ce modèle non entraîné:

 print("Input: \n", repr("".join(idx2char[input_example_batch[0]])))
print()
print("Next Char Predictions: \n", repr("".join(idx2char[sampled_indices ])))
 
Input: 
 "ader,\nTullus Aufidius, that will put you to 't.\nI sin in envying his nobility,\nAnd were I any thing "

Next Char Predictions: 
 "XEgmpNYR?'RUJEhE3uB;jZaxY.GKD;DW;p3S'dZLfRvrSgrfizpNCPTfZvaCQYcbB&OrBpg-BPM'V CDG,yS;,K$iUupH&m?WH Y"

Former le modèle

À ce stade, le problème peut être traité comme un problème de classification standard. Compte tenu de l'état RNN précédent et de l'entrée à ce pas de temps, prédisez la classe du caractère suivant.

Joindre un optimiseur et une fonction de perte

La fonction de perte standard tf.keras.losses.sparse_categorical_crossentropy fonctionne dans ce cas car elle est appliquée à travers la dernière dimension des prédictions.

Comme notre modèle renvoie des logits, nous devons définir l'indicateur from_logits .

 def loss(labels, logits):
  return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)

example_batch_loss  = loss(target_example_batch, example_batch_predictions)
print("Prediction shape: ", example_batch_predictions.shape, " # (batch_size, sequence_length, vocab_size)")
print("scalar_loss:      ", example_batch_loss.numpy().mean())
 
Prediction shape:  (64, 100, 65)  # (batch_size, sequence_length, vocab_size)
scalar_loss:       4.173355

Configurez la procédure de formation à l'aide de la méthode tf.keras.Model.compile . Nous utiliserons tf.keras.optimizers.Adam avec les arguments par défaut et la fonction de perte.

 model.compile(optimizer='adam', loss=loss)
 

Configurer les points de contrôle

Utilisez un tf.keras.callbacks.ModelCheckpoint pour vous assurer que les points de contrôle sont enregistrés pendant la formation:

 # Directory where the checkpoints will be saved
checkpoint_dir = './training_checkpoints'
# Name of the checkpoint files
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

checkpoint_callback=tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_prefix,
    save_weights_only=True)
 

Exécutez la formation

Pour que le temps de formation reste raisonnable, utilisez 10 époques pour entraîner le modèle. Dans Colab, définissez le runtime sur GPU pour un entraînement plus rapide.

 EPOCHS=10
 
 history = model.fit(dataset, epochs=EPOCHS, callbacks=[checkpoint_callback])
 
Epoch 1/10
172/172 [==============================] - 5s 26ms/step - loss: 2.7084
Epoch 2/10
172/172 [==============================] - 4s 26ms/step - loss: 1.9680
Epoch 3/10
172/172 [==============================] - 4s 26ms/step - loss: 1.6990
Epoch 4/10
172/172 [==============================] - 5s 27ms/step - loss: 1.5489
Epoch 5/10
172/172 [==============================] - 5s 27ms/step - loss: 1.4595
Epoch 6/10
172/172 [==============================] - 4s 26ms/step - loss: 1.3987
Epoch 7/10
172/172 [==============================] - 4s 26ms/step - loss: 1.3541
Epoch 8/10
172/172 [==============================] - 5s 26ms/step - loss: 1.3159
Epoch 9/10
172/172 [==============================] - 4s 26ms/step - loss: 1.2819
Epoch 10/10
172/172 [==============================] - 4s 26ms/step - loss: 1.2505

Générer du texte

Restaurer le dernier point de contrôle

Pour simplifier cette étape de prédiction, utilisez une taille de lot de 1.

En raison de la façon dont l'état RNN est passé d'un pas de temps à un autre, le modèle n'accepte qu'une taille de lot fixe une fois construit.

Pour exécuter le modèle avec un batch_size différent, nous devons reconstruire le modèle et restaurer les poids à partir du point de contrôle.

 tf.train.latest_checkpoint(checkpoint_dir)
 
'./training_checkpoints/ckpt_10'
 model = build_model(vocab_size, embedding_dim, rnn_units, batch_size=1)

model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))

model.build(tf.TensorShape([1, None]))
 
 model.summary()
 
Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding_1 (Embedding)      (1, None, 256)            16640     
_________________________________________________________________
gru_1 (GRU)                  (1, None, 1024)           3938304   
_________________________________________________________________
dense_1 (Dense)              (1, None, 65)             66625     
=================================================================
Total params: 4,021,569
Trainable params: 4,021,569
Non-trainable params: 0
_________________________________________________________________

La boucle de prédiction

Le bloc de code suivant génère le texte:

  • Il démarre en choisissant une chaîne de début, en initialisant l'état RNN et en définissant le nombre de caractères à générer.

  • Obtenez la distribution de prédiction du caractère suivant à l'aide de la chaîne de début et de l'état RNN.

  • Ensuite, utilisez une distribution catégorielle pour calculer l'indice du caractère prédit. Utilisez ce caractère prédit comme notre prochaine entrée dans le modèle.

  • L'état RNN renvoyé par le modèle est renvoyé dans le modèle afin qu'il ait maintenant plus de contexte, au lieu d'un seul caractère. Après avoir prédit le caractère suivant, les états RNN modifiés sont à nouveau réinjectés dans le modèle, c'est ainsi qu'il apprend à mesure qu'il obtient plus de contexte à partir des caractères prédits précédemment.

Pour générer du texte, la sortie du modèle est renvoyée à l'entrée

En regardant le texte généré, vous verrez que le modèle sait quand mettre en majuscule, créer des paragraphes et imiter un vocabulaire d'écriture semblable à celui de Shakespeare. Avec le petit nombre d'époques d'apprentissage, il n'a pas encore appris à former des phrases cohérentes.

 def generate_text(model, start_string):
  # Evaluation step (generating text using the learned model)

  # Number of characters to generate
  num_generate = 1000

  # Converting our start string to numbers (vectorizing)
  input_eval = [char2idx[s] for s in start_string]
  input_eval = tf.expand_dims(input_eval, 0)

  # Empty string to store our results
  text_generated = []

  # Low temperatures results in more predictable text.
  # Higher temperatures results in more surprising text.
  # Experiment to find the best setting.
  temperature = 1.0

  # Here batch size == 1
  model.reset_states()
  for i in range(num_generate):
      predictions = model(input_eval)
      # remove the batch dimension
      predictions = tf.squeeze(predictions, 0)

      # using a categorical distribution to predict the character returned by the model
      predictions = predictions / temperature
      predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()

      # We pass the predicted character as the next input to the model
      # along with the previous hidden state
      input_eval = tf.expand_dims([predicted_id], 0)

      text_generated.append(idx2char[predicted_id])

  return (start_string + ''.join(text_generated))
 
 print(generate_text(model, start_string=u"ROMEO: "))
 
ROMEO: there's now we heard of men,
One that way she hath to unluking bawd.'

GREMIO:
O, muthon, that lives us in noble limbs:
Like nothing?

BIANCA:
Ha! fare he had henced to do my brother.

ROMEO:
O Anuman!
O heavy, Pitizen:
Yes, I and therefore dadest I water?

ANGELO:
Yea, and air my honour, never saved in but self's grows encounty thousperous.

CLARENCE:
And therefore it is no hamp; your trankful mad
Whose ears bearned the moral ant intonson,
With what of death-man's estainted
By any object be sent--vice of death!
Come, you have deser'd at itself?

GENRD:
A pleasure may be tenger to me to him,
So stuties are.

PETRUCHIO:
'That, doth stood it
To ment so presently.
Thou hadst made guess that I should seem
I'll hanging, satisfied. Deny flore with his greatness!

FRIAR LAURENCE:
He lovedmy conveying so then toop allow:
Thy temper lost.

QUEEN:
'Tis set him him: hence, become an anot,
Bohemia will?

ISABELLA:
Yes;
Of bereff, you prithee so.

KING RICHARD III:
What sayest thou whilest venomed 

La chose la plus simple que vous puissiez faire pour améliorer les résultats est de l'entraîner plus longtemps (essayez EPOCHS=30 ).

Vous pouvez également expérimenter avec une chaîne de début différente, ou essayer d'ajouter une autre couche RNN pour améliorer la précision du modèle, ou ajuster le paramètre de température pour générer des prédictions plus ou moins aléatoires.

Avancé: formation personnalisée

La procédure de formation ci-dessus est simple, mais ne vous donne pas beaucoup de contrôle.

Maintenant que vous avez vu comment exécuter le modèle manuellement, décompressons la boucle d'entraînement et implémentons-la nous-mêmes. Cela donne un point de départ si, par exemple, pour mettre en œuvre l' apprentissage du curriculum pour aider à stabiliser la sortie en boucle ouverte du modèle.

Nous utiliserons tf.GradientTape pour suivre les dégradés. Vous pouvez en savoir plus sur cette approche en lisant le guide d'exécution impatient .

La procédure fonctionne comme suit:

  • Tout d'abord, réinitialisez l'état RNN. Nous faisons cela en appelant la méthode tf.keras.Model.reset_states .

  • Ensuite, parcourez l'ensemble de données (lot par lot) et calculez les prédictions associées à chacun.

  • Ouvrez un tf.GradientTape et calculez les prédictions et la perte dans ce contexte.

  • Calculez les gradients de la perte par rapport aux variables du modèle à l'aide de la méthode tf.GradientTape.grads .

  • Enfin, tf.train.Optimizer.apply_gradients un pas vers le bas en utilisant la méthode tf.train.Optimizer.apply_gradients l'optimiseur.

 model = build_model(
  vocab_size = len(vocab),
  embedding_dim=embedding_dim,
  rnn_units=rnn_units,
  batch_size=BATCH_SIZE)
 
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.iter
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.beta_1
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.beta_2
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.decay
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.learning_rate
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'm' for (root).layer_with_weights-0.embeddings
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'm' for (root).layer_with_weights-2.kernel
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'm' for (root).layer_with_weights-2.bias
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'm' for (root).layer_with_weights-1.cell.kernel
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'm' for (root).layer_with_weights-1.cell.recurrent_kernel
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'm' for (root).layer_with_weights-1.cell.bias
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'v' for (root).layer_with_weights-0.embeddings
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'v' for (root).layer_with_weights-2.kernel
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'v' for (root).layer_with_weights-2.bias
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'v' for (root).layer_with_weights-1.cell.kernel
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'v' for (root).layer_with_weights-1.cell.recurrent_kernel
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer's state 'v' for (root).layer_with_weights-1.cell.bias
WARNING:tensorflow:A checkpoint was restored (e.g. tf.train.Checkpoint.restore or tf.keras.Model.load_weights) but not all checkpointed values were used. See above for specific issues. Use expect_partial() on the load status object, e.g. tf.train.Checkpoint.restore(...).expect_partial(), to silence these warnings, or use assert_consumed() to make the check explicit. See https://www.tensorflow.org/guide/checkpoint#loading_mechanics for details.

 optimizer = tf.keras.optimizers.Adam()
 
 @tf.function
def train_step(inp, target):
  with tf.GradientTape() as tape:
    predictions = model(inp)
    loss = tf.reduce_mean(
        tf.keras.losses.sparse_categorical_crossentropy(
            target, predictions, from_logits=True))
  grads = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(grads, model.trainable_variables))

  return loss
 
 # Training step
EPOCHS = 10

for epoch in range(EPOCHS):
  start = time.time()

  # resetting the hidden state at the start of every epoch
  model.reset_states()

  for (batch_n, (inp, target)) in enumerate(dataset):
    loss = train_step(inp, target)

    if batch_n % 100 == 0:
      template = 'Epoch {} Batch {} Loss {}'
      print(template.format(epoch+1, batch_n, loss))

  # saving (checkpoint) the model every 5 epochs
  if (epoch + 1) % 5 == 0:
    model.save_weights(checkpoint_prefix.format(epoch=epoch))

  print ('Epoch {} Loss {:.4f}'.format(epoch+1, loss))
  print ('Time taken for 1 epoch {} sec\n'.format(time.time() - start))

model.save_weights(checkpoint_prefix.format(epoch=epoch))
 
Epoch 1 Batch 0 Loss 4.175529479980469
Epoch 1 Batch 100 Loss 2.3484625816345215
Epoch 1 Loss 2.1891
Time taken for 1 epoch 6.311128854751587 sec

Epoch 2 Batch 0 Loss 2.1668317317962646
Epoch 2 Batch 100 Loss 1.867101788520813
Epoch 2 Loss 1.7912
Time taken for 1 epoch 5.3653342723846436 sec

Epoch 3 Batch 0 Loss 1.8362562656402588
Epoch 3 Batch 100 Loss 1.6639163494110107
Epoch 3 Loss 1.6553
Time taken for 1 epoch 5.252068281173706 sec

Epoch 4 Batch 0 Loss 1.5616859197616577
Epoch 4 Batch 100 Loss 1.5669431686401367
Epoch 4 Loss 1.5058
Time taken for 1 epoch 5.2403037548065186 sec

Epoch 5 Batch 0 Loss 1.4984939098358154
Epoch 5 Batch 100 Loss 1.4885737895965576
Epoch 5 Loss 1.4406
Time taken for 1 epoch 5.334592819213867 sec

Epoch 6 Batch 0 Loss 1.3836933374404907
Epoch 6 Batch 100 Loss 1.4152636528015137
Epoch 6 Loss 1.4195
Time taken for 1 epoch 5.246667385101318 sec

Epoch 7 Batch 0 Loss 1.3539446592330933
Epoch 7 Batch 100 Loss 1.3515733480453491
Epoch 7 Loss 1.3782
Time taken for 1 epoch 5.239287376403809 sec

Epoch 8 Batch 0 Loss 1.2740596532821655
Epoch 8 Batch 100 Loss 1.291770577430725
Epoch 8 Loss 1.3270
Time taken for 1 epoch 5.19817590713501 sec

Epoch 9 Batch 0 Loss 1.2735955715179443
Epoch 9 Batch 100 Loss 1.238691806793213
Epoch 9 Loss 1.2839
Time taken for 1 epoch 5.2849626541137695 sec

Epoch 10 Batch 0 Loss 1.2129250764846802
Epoch 10 Batch 100 Loss 1.2772572040557861
Epoch 10 Loss 1.2674
Time taken for 1 epoch 5.498290538787842 sec