Учим нейронку писать как Фёдор Михайлович Достоевский

Сейчас мы обучим рекуррентную нейронную сеть создавать тексты в стиле Фёдора Михайловича Достоевского. Всё что ей для этого понадобится — это способность предсказывать следующую букву для строки из уже имеющихся. Не стоит ожидать от сети осмысленных фраз и предложений, но правила композиции слов, общую структуру и настроение она улавливает довольно неплохо.

Приведенный здесь результат работы получен за примерно 1 час обучения на ПК с видеокартой. Результат можно улучшить, увеличив время обучения или размер сети (её глубину или ширину слоев).

Итак, приступим!

Писать нейросеть мы будем на python, сейчас это фактически основной язык для Data Scientist. Использовать будем популярный фреймворк Keras, который позволяет очень просто описывать структуру нейронной сети и абстрагироваться от деталей её реализации. Keras внутри себя может использовать для вычислений библиотеки Tensorflow от Google или Theano. В нашем случае это будет Tensorflow. Библиотека поддерживает расчеты на GPU, так что мощная видеокарта от NVidia может ощутимо сократить время работы.

Загружаем все необходимые библиотеки, разрешаем бекенду Tensorflow увеличивать размер используемой GPU памяти при необходимости

Ввод:

 import tensorflow as tf
 from keras import backend as K
 config = tf.ConfigProto()
 config.gpu_options.allow_growth=True
 sess = tf.Session(config=config)
 K.set_session(sess) # разрешаем использовать больше видеопамяти
 from keras.models import Sequential
 from keras.layers import Dense, Activation, Dropout, LSTM,TimeDistributed, Embedding
 from keras.callbacks import ModelCheckpoint, Callback, EarlyStopping, ReduceLROnPlateau
 from keras.optimizers import RMSprop, Adam, SGD
 import sys
 import numpy as np

Using TensorFlow backend.

Загружаем текст, на котором будем обучаться и смотрим на его длину. Для более-менее интересного результата сети нужен передать на обучение больше миллиона символов текста.

Ввод:

fname = 'dostoevsky.txt' #тексты основных романов Федора Михайловича, объединенные в один файл
text = open(fname, 'r', encoding='utf-8').read()
print('corpus length:', len(text))

corpus length: 9469628

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

1. Число уникальных букв
2. Число уникальных пар букв (биграмм)
3. Число уникальных троек букв (триграмм)
4. Общее число слов и число уникальных слов

Для простоты дальше мы будем обучать нашу модель только на отдельных буквах.

Ввод:

chars = sorted(list(set(text)))
print('unique chars:', len(chars))
char_idx = dict((c, i) for i, c in enumerate(chars))
idx_char = dict((i, c) for i, c in enumerate(chars))

bigrams = []
for i in range(0, len(text)-1, 1):
   bigrams.append(text[i] + text[i+1])
bigrams = sorted(list(set(bigrams)))
print('unique bigrams:', len(bigrams))
bigram_idx = dict((c, i) for i, c in enumerate(bigrams))
idx_bigram = dict((i, c) for i, c in enumerate(bigrams))

trigrams = []
for i in range(0, len(text)-2, 1):
   trigrams.append(text[i] + text[i+1] + text[i+2])
trigrams = sorted(list(set(trigrams)))
print('unique trigrams:', len(trigrams))
trigram_idx = dict((c, i) for i, c in enumerate(trigrams))
idx_trigram = dict((i, c) for i, c in enumerate(trigrams))

words = text.split()
print('word count:', len(words))
words = sorted(list(set(words)))
print('unique words:', len(words))

unique chars: 162

unique bigrams: 3566 

unique trigrams: 26082 

word count: 1528727 

unique words: 165930

Ввод:

batch_size = 1024 # можно уменьшить, если при запуску возникает ошибка аллокации памяти
track_size = len(text) // batch_size
tracks = ['' for i in range(batch_size)]
for i in range(0, track_size):
   for track in range(batch_size):
       tracks[track] += text[track * track_size + i]

#  посмотрим что у нас получилось внутри отдельноно трека
print(tracks[4][:1000])

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

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

Ввод:

gram_idx = [char_idx, bigram_idx, trigram_idx] 
idx_gram = [idx_char, idx_bigram, idx_trigram]

# конвертируем блоки текста в наборы индексов символов
def grams(tracks, n=1):
   indexed = []
   for t in tracks:  
       track = []
       for i in range(0, len(t)-n+1, n):
           gram = ''
           for j in range(n):
               gram += t[i+j] # склеиваем каждые n символов в последовательность, получая символ, биграмму или триграмму в зависимости от n
           idx = gram_idx[n-1][gram] # ищем индекс последовательности в соответствующем словаре
           track.append(idx) # добавляем этот индекс в результирующий список
       indexed.append(track)
   return indexed

indexed = grams(tracks, 1)
vocab_size = len(gram_idx[0])
print('Vocabulary done: ', vocab_size)

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

Наша сеть будет предсказывать следующую букву для последовательности. Пример с фразой HELLO:

• мы начинаем подавать в сеть буквы одну за другой с начала. Одновременно на каждом шаге мы проверяем, что на выходе сети появилась следующая буква из этого слова.
• вначале подаем на вход H, на выходе ждем E
• на следующем шаге подаем E, ждём на выходе L
• подаем L, ждём L
• подаем L, ждём O.

Всё это время сеть сохраняет своё состояние с самого начала, что позволяет ей запомнить, в какой ситуации после L надо ответить L, а в какой O.

Для предсказания индекса буквы нам необходимо преобразовать его в one-hot вектор — длина которого равна размеру нашего словаря, а во всех позициях кроме позиции с номером текущего индекса стоят нули. В позиции индекса же будет стоять единица. Например, для словаря из [E, H, L, O] длина словаря была бы 4, индекс буквы L — 3, а её one-hot вектор 0010.

Также мы расставляем блоки из последовательностей длины N так, чтобы в каждом новом батче на позиции 1 стояла последовательность, продолжающая последовательность 1 из предыдущего батча, на позиции 2 — продолжающая текст позиции 2 предыдущего батча, и так далее.

Ввод:

seq_len = 40 # длина последовательности в батче

def onehot(n):
   v = [0 for i in range(vocab_size)]
   v[n] = 1
   return v

# расставляем последовательности из блоков таким образом, что первая последовательность из следующего батча продолжает первую последовательность текущего. Чтобы слой LSTM мог использовать состояние из предыдущего батча
def vectorize(tracks):
   track_size = len(tracks[0])
   X = []
   y = []
   for i in range(0, track_size - seq_len + 1, seq_len):
       for t in tracks:
           X.append(t[i:i + seq_len])
           target = [onehot(c) for c in t[i+1:i + seq_len + 1]]
           y.append(target)
   return X, y
       
x,y = vectorize(indexed)
print('Number of training samples', len(x))
print('Number of training labels', len(y))
print('Label sequence length', len(y[0]))
print('Label character one-hot vector length', len(y[0][0]))

Number of training samples 236544
Number of training labels 236544
Label sequence length 40
Label character one-hot vector length 162

Приготовим нашу рекуррентную сеть. В начале идёт embedding-слой, преобразующий индексы входной последовательности в dense вектор фиксированного размера. Затем несколько рекуррентных LSTM слоёв и один полносвязный слой с число выходов, равным размеру нашего словаря. В качестве loss-функции нам нужна кросс энтропия, в качестве оптимизатора будем использовать популярный сейчас алгоритм Adam, также в нём важно не забыть установить clipnorm — ограничение на размер градиентов.

Ввод:

cells = 512
drop = 0.2
embed = 30 # размер вектора символа (эмбеддинга)
layers = 2
lr = 0.01
clip = 5.0 # ограничение градиентов при оптимизации
stateful=True # сохраняем состояние слоя между батчами

model = Sequential()
model.add(Embedding(vocab_size, embed, batch_input_shape=(batch_size, seq_len)))
for l in range(layers):
   model.add(LSTM(cells, return_sequences=True, stateful=stateful, dropout=drop))
model.add(Dense(vocab_size))
model.add(Activation('softmax'))
optimizer = Adam(lr, clipnorm=clip)
model.compile(loss="categorical_crossentropy", optimizer=optimizer)
model.summary()

Как видно, размер сети получился порядка 3.3 миллиона параметров.

Ввод:

def get_callbacks(filepath, patience=5):
   learning_rate_reduction = ReduceLROnPlateau(monitor='loss', 
                                           patience=patience, 
                                           verbose=1, 
                                           factor=0.5, 
                                           min_lr=0.00001)
   es = EarlyStopping('loss', verbose=1, min_delta=0.02, patience=patience, mode="min")
   return [learning_rate_reduction, es]

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

Ввод:

def sample(preds, temperature=0.5):
   preds = np.asarray(preds).astype('float64')
   preds = np.log(preds) / temperature
   exp_preds = np.exp(preds)
   preds = exp_preds / np.sum(exp_preds)
   probas = np.random.multinomial(1, preds, 1)
   return np.argmax(probas)

Наша изначальная модель была построена исходя из указанного размера батча, а значит плохо подходит для работы на одной последовательности. Создадим её копию с размером батча 1 и перенесем туда веса обученной модели
Ввод:

def predictive_model(main_model):

   # меняем размер батча на 1, чтобы подавать всего один блок текста
   model = Sequential()
   model.add(Embedding(vocab_size, embed, batch_input_shape=(1, seq_len)))
   for l in range(layers):
       model.add(LSTM(cells, return_sequences=True, stateful=stateful, dropout=drop))
   model.add(Dense(vocab_size))
   model.add(Activation('softmax'))
   optimizer = Adam(lr, clipnorm=clip)
   model.compile(loss="categorical_crossentropy", optimizer=optimizer)
   old_weights = main_model.get_weights()
   model.set_weights(old_weights)
   return model

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

# генерируем новый символ по имеющимся с помощью нашей модели, добавляем его к строке и опять генерируем уже по этой дополненной строке
def test(model, l=500, seed = None, t=0.5):
   start_from = np.random.randint(len(text)-seq_len)+seq_len
   seed_string = text[start_from:start_from + seq_len*3] if seed is None else seed
   print('\n\nSeed:  ', seed_string)
   print('----')
   sys.stdout.write(seed_string)
   prmodel = predictive_model(model)
   for i in range(l):
       prmodel.reset_states()
       padlen = (len(seed_string) // seq_len +1) * seq_len
       seed_string = seed_string.rjust(padlen)[-seq_len*3:]
       test_tracks = [seed_string]
       tidx = grams(test_tracks)
       xt, _ = vectorize(tidx)
       preds = prmodel.predict(np.array(xt), batch_size=1, verbose=0)
       preds = preds[-1][-1] # last symbol of last sequence
       next_item = idx_char[sample(preds, t)]
       seed_string = seed_string + next_item
       sys.stdout.write(next_item)
       sys.stdout.flush()

Пришло время обучить сеть. Попробуем делать это на протяжении 50 итераций, периодически (раз в 3 итерации) оценивая качество сгенерированного сетью текста

Ввод:

for iteration in range(1, 51):
   print('\nIteration', iteration)
   model_name = 'char_%s_%d_%d_%.1f_%d.h5' % (fname, layers, cells, drop, iteration)
   history=model.fit(
       np.array(x), np.array(y), 
       batch_size=batch_size, 
       epochs=1, 
       verbose=1, 
       shuffle=False,
       callbacks=get_callbacks(filepath=model_name)
   )
   model.save_weights(model_name, overwrite=True)
   model.reset_states()
   if iteration%3 == 0:
       test(model)

Теперь мы можем посмотреть как влияет температура сэмплинга на выводимый моделью текст. Используем для этого одну и ту же начальную последовательность и посмотрим на результаты

Ввод:

seed = "У нас в Малом зале до сих пор проходят"
test(model, 300, seed, 0.1)
test(model, 300, seed, 0.5)
test(model, 300, seed, 0.7)

Автор: Олег Шляжко

Ссылка на github: http://github.com/ollmer/neural_experiments/blob/master/char_rnn_dostoevsky.ipynb