Эта страница была переведа с помощью Cloud Translation API.
Switch to English

Генерация текста с помощью RNN

Посмотреть на TensorFlow.org Запускаем в Google Colab Посмотреть исходный код на GitHub Скачать блокнот

В этом руководстве показано, как сгенерировать текст с помощью символьной RNN. Вы будете работать с набором данных из сочинений Шекспира из книги Андрея Карпати «Неоправданная эффективность рекуррентных нейронных сетей» . Учитывая последовательность символов из этих данных («Шекспир»), обучите модель предсказывать следующий символ в последовательности («е»). Более длинные последовательности текста можно сгенерировать, повторно вызывая модель.

Это руководство включает исполняемый код, реализованный с использованием tf.keras и активного выполнения . Ниже приводится пример выходных данных, когда модель в этом руководстве обучалась в течение 30 эпох и начиналась со строки «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

Хотя некоторые предложения грамматичны, большинство из них не имеет смысла. Модель не выучила значения слов, но учтите:

  • Модель основана на персонажах. Когда началось обучение, модель не знала, как написать английское слово или что слова даже были единицей текста.

  • Структура вывода напоминает пьесу - блоки текста обычно начинаются с имени докладчика заглавными буквами, как и в наборе данных.

  • Как показано ниже, модель обучается на небольших пакетах текста (по 100 символов каждый) и по-прежнему способна генерировать более длинную последовательность текста с согласованной структурой.

Настроить

Импортировать TensorFlow и другие библиотеки

import tensorflow as tf

import numpy as np
import os
import time

Скачать набор данных Шекспира

Измените следующую строку, чтобы запустить этот код на ваших собственных данных.

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

Прочтите данные

Сначала загляните в текст:

# 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

Обработать текст

Векторизовать текст

Перед обучением вам необходимо отобразить строки в числовое представление. Создайте две таблицы поиска: одна отображает символы в числа, а другая - для чисел в символы.

# 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])

Теперь у вас есть целочисленное представление для каждого символа. Обратите внимание, что вы сопоставили символ как индексы от 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]

Задача прогнозирования

Для данного символа или последовательности символов какой следующий символ наиболее вероятен? Это задача, которой вы обучаете модель. Входными данными для модели будет последовательность символов, и вы обучаете модель предсказывать выходные данные - следующий символ на каждом временном шаге.

Поскольку RNN поддерживают внутреннее состояние, которое зависит от ранее замеченных элементов, учитывая все символы, вычисленные до этого момента, какой будет следующий символ?

Создавайте обучающие примеры и цели

Затем разделите текст на примеры последовательностей. Каждая входная последовательность будет содержать символы seq_length из текста.

Для каждой входной последовательности соответствующие целевые объекты содержат текст одинаковой длины, за исключением смещения на один символ вправо.

Так что разбейте текст на куски seq_length+1 . Например, предположим, что seq_length равно 4, а наш текст - «Hello». Входной последовательностью будет «Ад», а целевой последовательностью - «Привет».

Для этого сначала используйте функцию tf.data.Dataset.from_tensor_slices чтобы преобразовать текстовый вектор в поток индексов символов.

# The maximum length sentence you 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

batch метод позволяет нам легко преобразовать эти отдельные символы в последовательности желаемого размера.

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'

Для каждой последовательности продублируйте и сдвиньте ее, чтобы сформировать входной и целевой текст, используя метод map для применения простой функции к каждому пакету:

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

dataset = sequences.map(split_input_target)

Распечатайте первый пример входных и целевых значений:

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 '

Каждый индекс этих векторов обрабатывается как один временной шаг. Для ввода на временном шаге 0 модель получает индекс для «F» и пытается предсказать индекс для «i» в качестве следующего символа. На следующем временном шаге он делает то же самое, но RNN учитывает контекст предыдущего шага в дополнение к текущему входному символу.

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 (' ')

Создавайте тренировочные партии

Вы использовали tf.data чтобы разбить текст на управляемые последовательности. Но прежде чем вводить эти данные в модель, вам нужно перетасовать данные и упаковать их в пакеты.

# 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)>

Постройте модель

Используйте tf.keras.Sequential для определения модели. В этом простом примере для определения нашей модели используются три слоя:

  • tf.keras.layers.Embedding : входной слой. Обучаемая справочная таблица, которая отображает номера каждого символа в вектор с размерами embedding_dim ;
  • tf.keras.layers.GRU : Тип RNN с размером units=rnn_units (здесь также можно использовать слой LSTM.)
  • tf.keras.layers.Dense : выходной слой с выходными значениями 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)

Для каждого символа модель ищет встраивание, запускает GRU за один шаг с встраиванием в качестве входных данных и применяет плотный слой для генерации логитов, прогнозирующих логарифмическую вероятность следующего символа:

Чертеж данных, проходящих через модель

Обратите внимание, что здесь используется последовательная модель Keras, поскольку все слои модели имеют только один вход и производят один выход. Если вы хотите получить и повторно использовать состояния из уровня RNN с отслеживанием состояния, вы можете создать свою модель с помощью функционального API Keras или создания подклассов модели. Пожалуйста, обратитесь к руководству Keras RNN для получения более подробной информации.

Попробуйте модель

Теперь запустите модель, чтобы убедиться, что она ведет себя должным образом.

Сначала проверьте форму вывода:

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)

В приведенном выше примере длина последовательности входных данных равна 100 но модель может быть запущена на входах любой длины:

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
_________________________________________________________________

Чтобы получить фактические прогнозы от модели, вам необходимо выполнить выборку из выходного распределения, чтобы получить фактические индексы символов. Это распределение определяется логитами по словарю символов.

Попробуйте это для первого примера в партии:

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

Это дает нам на каждом временном шаге прогноз следующего индекса символа:

sampled_indices
array([41, 60,  3, 31, 47, 21, 61,  6, 56, 42, 39, 40, 52, 60, 37, 37, 27,
       11,  6, 56, 64, 62, 43, 42,  6, 34,  1, 30, 16, 45, 46, 11, 17,  8,
       26,  8,  1, 46, 37, 21, 37, 53, 34, 49,  5, 58, 11,  9, 42, 62, 14,
       56, 56, 30, 31, 32, 63, 53, 10, 23, 35,  5, 19, 19, 46,  3, 23, 63,
       61, 11, 57,  0, 35, 48, 32,  4, 37,  7, 48, 23, 39, 30, 20, 26,  1,
       52, 57, 23, 46, 56, 11, 22,  7, 47, 16, 27, 38, 51, 55, 28])

Расшифруйте их, чтобы увидеть текст, предсказанный этой неподготовленной моделью:

print("Input: \n", repr("".join(idx2char[input_example_batch[0]])))
print()
print("Next Char Predictions: \n", repr("".join(idx2char[sampled_indices ])))
Input: 
 "dness! Make not impossible\nThat which but seems unlike: 'tis not impossible\nBut one, the wicked'st c"

Next Char Predictions: 
 "cv$SiIw,rdabnvYYO;,rzxed,V RDgh;E.N. hYIYoVk't;3dxBrrRSTyo:KW'GGh$Kyw;s\nWjT&Y-jKaRHN nsKhr;J-iDOZmqP"

Обучите модель

На этом этапе проблему можно рассматривать как стандартную задачу классификации. Учитывая предыдущее состояние RNN и ввод этого временного шага, спрогнозируйте класс следующего символа.

Прикрепите оптимизатор и функцию потерь

Стандартная tf.keras.losses.sparse_categorical_crossentropy потерь tf.keras.losses.sparse_categorical_crossentropy работает в этом случае, потому что она применяется ко всему последнему измерению прогнозов.

Поскольку ваша модель возвращает логиты, вам необходимо установить флаг 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.174373

Настройте процедуру обучения с tf.keras.Model.compile метода tf.keras.Model.compile . Используйте tf.keras.optimizers.Adam с аргументами по умолчанию и функцией потерь.

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

Настроить контрольные точки

Используйте tf.keras.callbacks.ModelCheckpoint чтобы убедиться, что контрольные точки сохраняются во время обучения:

# 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)

Провести обучение

Чтобы время обучения было разумным, используйте 10 эпох для обучения модели. В Colab установите время выполнения на GPU для более быстрого обучения.

EPOCHS = 10
history = model.fit(dataset, epochs=EPOCHS, callbacks=[checkpoint_callback])
Epoch 1/10
172/172 [==============================] - 5s 27ms/step - loss: 2.6807
Epoch 2/10
172/172 [==============================] - 5s 27ms/step - loss: 1.9748
Epoch 3/10
172/172 [==============================] - 5s 26ms/step - loss: 1.7063
Epoch 4/10
172/172 [==============================] - 5s 26ms/step - loss: 1.5543
Epoch 5/10
172/172 [==============================] - 5s 27ms/step - loss: 1.4633
Epoch 6/10
172/172 [==============================] - 5s 26ms/step - loss: 1.4028
Epoch 7/10
172/172 [==============================] - 5s 26ms/step - loss: 1.3568
Epoch 8/10
172/172 [==============================] - 5s 26ms/step - loss: 1.3187
Epoch 9/10
172/172 [==============================] - 5s 26ms/step - loss: 1.2845
Epoch 10/10
172/172 [==============================] - 5s 26ms/step - loss: 1.2528

Создать текст

Восстановить последнюю контрольную точку

Чтобы упростить этот шаг прогнозирования, используйте размер пакета 1.

Из-за способа передачи состояния RNN от временного шага к временному шагу модель принимает только фиксированный размер пакета после построения.

Чтобы запустить модель с другим batch_size , вам необходимо перестроить модель и восстановить веса из контрольной точки.

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
_________________________________________________________________

Цикл предсказания

Следующий блок кода генерирует текст:

  • Начните с выбора начальной строки, инициализации состояния RNN и установки количества символов для генерации.

  • Получите прогнозируемое распределение следующего символа, используя начальную строку и состояние RNN.

  • Затем используйте категориальное распределение, чтобы вычислить индекс предсказанного символа. Используйте этот предсказанный символ как наш следующий вход в модель.

  • Состояние RNN, возвращаемое моделью, передается обратно в модель, так что теперь у нее больше контекста, а не только один символ. После предсказания следующего символа измененные состояния RNN снова передаются обратно в модель, и именно так она обучается, поскольку получает больше контекста от ранее предсказанных символов.

Чтобы сгенерировать текст, выход модели возвращается на вход

Глядя на сгенерированный текст, вы увидите, что модель знает, когда писать с заглавной буквы, создавать абзацы и имитирует словарный запас, подобный шекспировскому письму. При небольшом количестве эпох обучения он еще не научился составлять связные предложения.

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 temperature results in more predictable text.
    # Higher temperature 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()

        # 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: ghast I cut go,
Know the normander and the wrong:
To our Morsuis misdress are behiod;
And after as if no other husion.

VALERIS:
Your father and of worms?

LADY GREY:
Your hot can dost.

WARWICK:
Then, atient the bade, truckle aid,
Dearve your tongue should be cred to our face,
Bear trouble my father valiant,' in the company.

SICINIUS:
O God!'Sir afeard?

MIRANDA:
Come, good med,---or whom by the duke?

DUKE VINCENTIO:
Yes, that are bore indocation!

IO:
None not, my lord's sons.

MIRANDA:
Of some King?'
And, if thou was, a partanot young to thee.

JULIET:
O, tell; then I'll see them again? There's not so reder
no mother, and my three here to us. You might shall not speak, these this
same this within; what armpy I might
but though some way.

ROMEO:
Our daughter of the fool, that great come.
So, not the sun summer so all the sends,
Your ludgers made before the souls of years, and thereby there. Lady, father, were well the sold, pass, remeate.

Second King Richard's daughter,
Which chee

Самое простое, что вы можете сделать для улучшения результатов, - это тренировать его дольше (попробуйте EPOCHS = 30 ).

Вы также можете поэкспериментировать с другой начальной строкой, попробовать добавить еще один слой RNN, чтобы повысить точность модели, или отрегулировать параметр температуры для генерации более или менее случайных прогнозов.

Продвинутый: индивидуальное обучение

Вышеуказанная процедура обучения проста, но не дает вам особого контроля.

Итак, теперь, когда вы увидели, как запустить модель вручную, давайте распакуем цикл обучения и реализуем его сами. Это дает отправную точку, если, например, вы хотите реализовать обучение по учебной программе, чтобы помочь стабилизировать выход модели без обратной связи.

Используйте tf.GradientTape чтобы отслеживать градиенты. Вы можете узнать больше об этом подходе, прочитав руководство по быстрому выполнению .

Процедура работает следующим образом:

  • Сначала сбросьте состояние RNN. Вы делаете это, вызывая метод tf.keras.Model.reset_states .

  • Затем выполните итерацию по набору данных (пакет за пакетом) и вычислите прогнозы, связанные с каждым из них.

  • Откройте tf.GradientTape и рассчитайте прогнозы и потери в этом контексте.

  • Рассчитайте градиенты потерь по отношению к переменным модели, используя метод tf.GradientTape.grads .

  • Наконец, сделайте шаг вниз, используя метод оптимизатора tf.train.Optimizer.apply_gradients .

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.174976348876953
Epoch 1 Batch 100 Loss 2.351067304611206
Epoch 1 Loss 2.1421
Time taken for 1 epoch 6.3171796798706055 sec

Epoch 2 Batch 0 Loss 2.166642665863037
Epoch 2 Batch 100 Loss 1.9492360353469849
Epoch 2 Loss 1.7901
Time taken for 1 epoch 5.3413612842559814 sec

Epoch 3 Batch 0 Loss 1.804692029953003
Epoch 3 Batch 100 Loss 1.6545528173446655
Epoch 3 Loss 1.6328
Time taken for 1 epoch 5.337632179260254 sec

Epoch 4 Batch 0 Loss 1.6188888549804688
Epoch 4 Batch 100 Loss 1.5314372777938843
Epoch 4 Loss 1.5319
Time taken for 1 epoch 5.2844321727752686 sec

Epoch 5 Batch 0 Loss 1.470827579498291
Epoch 5 Batch 100 Loss 1.4400928020477295
Epoch 5 Loss 1.4442
Time taken for 1 epoch 5.46646785736084 sec

Epoch 6 Batch 0 Loss 1.4113285541534424
Epoch 6 Batch 100 Loss 1.387071132659912
Epoch 6 Loss 1.3713
Time taken for 1 epoch 5.243147373199463 sec

Epoch 7 Batch 0 Loss 1.3486154079437256
Epoch 7 Batch 100 Loss 1.353363037109375
Epoch 7 Loss 1.3270
Time taken for 1 epoch 5.295132160186768 sec

Epoch 8 Batch 0 Loss 1.2960264682769775
Epoch 8 Batch 100 Loss 1.3038402795791626
Epoch 8 Loss 1.3556
Time taken for 1 epoch 5.228798151016235 sec

Epoch 9 Batch 0 Loss 1.2495232820510864
Epoch 9 Batch 100 Loss 1.30863618850708
Epoch 9 Loss 1.2699
Time taken for 1 epoch 5.33559775352478 sec

Epoch 10 Batch 0 Loss 1.2161246538162231
Epoch 10 Batch 100 Loss 1.2242770195007324
Epoch 10 Loss 1.2360
Time taken for 1 epoch 5.377742528915405 sec