AI Writing Poems: построение модели LSTM с использованием PyTorch

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

Что такое РНН?

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

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

Что такое LSTM?

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

Пример 1: «Птицы живут в гнезде». Здесь легко предугадать слово «гнездо», потому что у нас есть предыдущий контекст птицы, и в этом случае RNN будет работать нормально.

Пример 2:

«Я вырос в Индии… .. так что я могу говорить на х инди», поэтому здесь задача предсказания слова «хинди» трудна для RNN, потому что здесь большой разрыв между контекстами. Глядя на строку «Я могу говорить…», мы не можем предсказать, на каком языке нам понадобится дополнительный контекст Индии. Итак, здесь нам понадобится некоторая долгосрочная зависимость от нашего абзаца, чтобы мы могли понять контекст.

Для этого мы используем LSTM (Long Short Term Memory). Судя по названию, у них есть долговременная и кратковременная память (ворота), и обе они используются вместе, чтобы делать прогнозы. Если мы говорим об архитектуре LSTM, они содержат 4 шлюза, а именно: изучить ворота, забыть ворота, запомнить ворота, использовать ворота. Чтобы статья была простой и удобной, я не буду вдаваться в теорию архитектуры LSTM. Но, возможно, мы поговорим об этом в следующих статьях (может быть, в следующей 😉).

Давайте построим нашу модель.

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

Загрузка и предварительная обработка данных

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

1. Сначала импортируйте библиотеки.

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

2. Теперь загрузите данные из текстового файла.

# open text file and read in data as `text`
with open('/data/poems_data.txt', 'r') as f:
    text = f.read()

3. Мы можем проверить наши данные, напечатав первые 100 символов.

text[:100]

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

# encode the text and map each character to an integer and vice versa
# we create two dictionaries:
# 1. int2char, which maps integers to characters
# 2. char2int, which maps characters to unique integers
chars = tuple(set(text))
int2char = dict(enumerate(chars))
char2int = {ch: ii for ii, ch in int2char.items()}
# encode the text
encoded = np.array([char2int[ch] for ch in text])

5. Для представления символа используется одна горячая кодировка. например, если у нас есть три символа a, b, c, тогда мы можем представить их таким образом [1,0,0], [0,1,0], [0,0,1] здесь мы используем 1 для представления этого символа а все остальные будут равны 0. Для нашего случая использования у нас много символов и символов, поэтому наш горячий вектор будет длинным. Но это нормально.

def one_hot_encode(arr, n_labels):
    
    # Initialize the the encoded array
    one_hot = np.zeros((arr.size, n_labels), dtype=np.float32)
    
    # Fill the appropriate elements with ones
    one_hot[np.arange(one_hot.shape[0]), arr.flatten()] = 1.
    
    # Finally reshape it to get back to the original array
    one_hot = one_hot.reshape((*arr.shape, n_labels))
    
    return one_hot

Теперь проверьте это таким образом.

# check that the function works as expected
test_seq = np.array([[0, 5, 1]])
one_hot = one_hot_encode(test_seq, 8)
print(one_hot)

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

def get_batches(arr, batch_size, seq_length):
    '''Create a generator that returns batches of size
       batch_size x seq_length from arr.
       
       Arguments
       ---------
       arr: Array you want to make batches from
       batch_size: Batch size, the number of sequences per batch
       seq_length: Number of encoded chars in a sequence
    '''
    
    batch_size_total = batch_size * seq_length
    # total number of batches we can make
    n_batches = len(arr)//batch_size_total
    
    # Keep only enough characters to make full batches
    arr = arr[:n_batches * batch_size_total]
    # Reshape into batch_size rows
    arr = arr.reshape((batch_size, -1))
    
    # iterate through the array, one sequence at a time
    for n in range(0, arr.shape[1], seq_length):
        # The features
        x = arr[:, n:n+seq_length]
        # The targets, shifted by one
        y = np.zeros_like(x)
        try:
            y[:, :-1], y[:, -1] = x[:, 1:], arr[:, n+seq_length]
        except IndexError:
            y[:, :-1], y[:, -1] = x[:, 1:], arr[:, 0]
        yield x, y

7. теперь мы можем проверить, доступен ли графический процессор (если графический процессор недоступен, оставьте минимальное количество эпох)

# check if GPU is available
train_on_gpu = torch.cuda.is_available()
if(train_on_gpu):
    print('Training on GPU!')
else: 
    print('No GPU available, training on CPU; consider making n_epochs very small.')

8. Здесь мы создали класс CharRNN. Это наш класс модели. В методе init мы должны определить слои для нашей модели. Здесь мы используем два слоя LSTM. Мы также используем отсев (это помогает избежать переобучения). Для вывода мы использовали простой линейный слой.

class CharRNN(nn.Module):
    
    def __init__(self, tokens, n_hidden=256, n_layers=2,
                               drop_prob=0.5, lr=0.001):
        super().__init__()
        self.drop_prob = drop_prob
        self.n_layers = n_layers
        self.n_hidden = n_hidden
        self.lr = lr
        
        # creating character dictionaries
        self.chars = tokens
        self.int2char = dict(enumerate(self.chars))
        self.char2int = {ch: ii for ii, ch in self.int2char.items()}
        
        #lstm layer
        self.lstm=nn.LSTM(len(self.chars),n_hidden,n_layers,
                          dropout=drop_prob,batch_first=True)
        
        #dropout layer
        self.dropout=nn.Dropout(drop_prob)
        
        #output layer
        self.fc=nn.Linear(n_hidden,len(self.chars))
    
    def forward(self, x, hidden):
        ''' Forward pass through the network. 
            These inputs are x, and the hidden/cell state `hidden`. '''
        ## Get the outputs and the new hidden state from the lstm
        r_output, hidden = self.lstm(x, hidden)
        
        ## pass through a dropout layer
        out = self.dropout(r_output)
        
        # Stack up LSTM outputs using view
        # you may need to use contiguous to reshape the output
        out = out.contiguous().view(-1, self.n_hidden)
        
        ## put x through the fully-connected layer
        out = self.fc(out)
        return out, hidden
    
    
    def init_hidden(self, batch_size):
        ''' Initializes hidden state '''
        # Create two new tensors with sizes n_layers x batch_size x n_hidden,
        # initialized to zero, for hidden state and cell state of LSTM
        weight = next(self.parameters()).data
        
        if (train_on_gpu):
            hidden = (weight.new(self.n_layers, batch_size, self.n_hidden).zero_().cuda(),
                  weight.new(self.n_layers, batch_size, self.n_hidden).zero_().cuda())
        else:
            hidden = (weight.new(self.n_layers, batch_size, self.n_hidden).zero_(),
                      weight.new(self.n_layers, batch_size, self.n_hidden).zero_())
        
        return hidden

9. Теперь у нас есть наша модель, пора ее обучить. Для обучения мы должны использовать оптимизатор и функцию потерь. Мы просто вычисляем потери после каждого шага, затем шаговая функция оптимизатора распространяет их обратно, и значения изменяются соответствующим образом. Потери будут постепенно уменьшаться, а это значит, что модель улучшается.

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

def train(net, data, epochs=10, batch_size=10, seq_length=50, lr=0.001, clip=5, val_frac=0.1, print_every=10):
    ''' Training a network 
    
        Arguments
        ---------
        
        net: CharRNN network
        data: text data to train the network
        epochs: Number of epochs to train
        batch_size: Number of mini-sequences per mini-batch, aka batch size
        seq_length: Number of character steps per mini-batch
        lr: learning rate
        clip: gradient clipping
        val_frac: Fraction of data to hold out for validation
        print_every: Number of steps for printing training and validation loss
    
    '''
    net.train()
    
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    
    # create training and validation data
    val_idx = int(len(data)*(1-val_frac))
    data, val_data = data[:val_idx], data[val_idx:]
    
    if(train_on_gpu):
        net.cuda()
    
    counter = 0
    n_chars = len(net.chars)
    for e in range(epochs):
        # initialize hidden state
        h = net.init_hidden(batch_size)
        
        for x, y in get_batches(data, batch_size, seq_length):
            counter += 1
            
            # One-hot encode our data and make them Torch tensors
            x = one_hot_encode(x, n_chars)
            inputs, targets = torch.from_numpy(x), torch.from_numpy(y)
            
            if(train_on_gpu):
                inputs, targets = inputs.cuda(), targets.cuda()
            # Creating new variables for the hidden state, otherwise
            # we'd backprop through the entire training history
            h = tuple([each.data for each in h])
            # zero accumulated gradients
            net.zero_grad()
            
            # get the output from the model
            output, h = net(inputs, h)
            
            # calculate the loss and perform backprop
            loss = criterion(output, targets.view(batch_size*seq_length).long())
            loss.backward()
            # `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
            nn.utils.clip_grad_norm_(net.parameters(), clip)
            opt.step()
            
            # loss stats
            if counter % print_every == 0:
                # Get validation loss
                val_h = net.init_hidden(batch_size)
                val_losses = []
                net.eval()
                for x, y in get_batches(val_data, batch_size, seq_length):
                    # One-hot encode our data and make them Torch tensors
                    x = one_hot_encode(x, n_chars)
                    x, y = torch.from_numpy(x), torch.from_numpy(y)
                    
                    # Creating new variables for the hidden state, otherwise
                    # we'd backprop through the entire training history
                    val_h = tuple([each.data for each in val_h])
                    
                    inputs, targets = x, y
                    if(train_on_gpu):
                        inputs, targets = inputs.cuda(), targets.cuda()
                    output, val_h = net(inputs, val_h)
                    val_loss = criterion(output, targets.view(batch_size*seq_length).long())
                
                    val_losses.append(val_loss.item())
                
                net.train() # reset to train mode after iterationg through validation data
                
                print("Epoch: {}/{}...".format(e+1, epochs),
                      "Step: {}...".format(counter),
                      "Loss: {:.4f}...".format(loss.item()),
                      "Val Loss: {:.4f}".format(np.mean(val_losses)))

теперь тренируйте его следующим образом.

# define and print the net
n_hidden = 512
n_layers = 2
net = CharRNN(chars, n_hidden, n_layers)
print(net)
batch_size = 128
seq_length = 100
n_epochs =  10# start small if you are just testing initial behavior
# train the model
train(net, encoded, epochs=n_epochs, batch_size=batch_size, seq_length=seq_length, lr=0.001, print_every=10)

10. Сохраняем модель следующим образом.

# change the name, for saving multiple files
model_name = 'poem_4_epoch.net'
checkpoint = {'n_hidden': net.n_hidden,
              'n_layers': net.n_layers,
              'state_dict': net.state_dict(),
              'tokens': net.chars}
with open(model_name, 'wb') as f:
    torch.save(checkpoint, f)

11. Теперь, когда модель обучена, мы хотим сделать выборку из нее и сделать прогнозы относительно следующих персонажей! Для выборки мы передаем символ, и сеть предсказывает следующий символ. Затем мы берем этот символ, передаем его обратно и получаем еще один предсказанный символ. Просто продолжайте делать это, и вы создадите кучу текста! Здесь верхний k выборок - это просто номер буквы, которую наша модель будет предугадывать и использовать наиболее релевантную из них.

def predict(net, char, h=None, top_k=None):
        ''' Given a character, predict the next character.
            Returns the predicted character and the hidden state.
        '''
        
        # tensor inputs
        x = np.array([[net.char2int[char]]])
        x = one_hot_encode(x, len(net.chars))
        inputs = torch.from_numpy(x)
        
        if(train_on_gpu):
            inputs = inputs.cuda()
        
        # detach hidden state from history
        h = tuple([each.data for each in h])
        # get the output of the model
        out, h = net(inputs, h)
        # get the character probabilities
        p = F.softmax(out, dim=1).data
        if(train_on_gpu):
            p = p.cpu() # move to cpu
        
        # get top characters
        if top_k is None:
            top_ch = np.arange(len(net.chars))
        else:
            p, top_ch = p.topk(top_k)
            top_ch = top_ch.numpy().squeeze()
        
        # select the likely next character with some element of randomness
        p = p.numpy().squeeze()
        char = np.random.choice(top_ch, p=p/p.sum())
        
        # return the encoded value of the predicted char and the hidden state
        return net.int2char[char], h
def sample(net, size, prime='The', top_k=None):
        
    if(train_on_gpu):
        net.cuda()
    else:
        net.cpu()
    
    net.eval() # eval mode
    
    # First off, run through the prime characters
    chars = [ch for ch in prime]
    h = net.init_hidden(1)
    for ch in prime:
        char, h = predict(net, ch, h, top_k=top_k)
    chars.append(char)
    
    # Now pass in the previous character and get a new one
    for ii in range(size):
        char, h = predict(net, chars[-1], h, top_k=top_k)
        chars.append(char)
    return ''.join(chars)

12. Теперь мы можем использовать этот примерный метод для прогнозирования.

print(sample(net, 500, prime='christmas', top_k=2))

и вывод будет выглядеть примерно так.

christmas a son of this
the sun wants the street of the stars, and the way the way
they went and too man and the star of the words
of a body of a street and the strange shoulder of the sky
and the sun, an end on the sun and the sun and so to the stars are stars
and the words of the water and the streets of the world
to see them to start a posture of the streets
on the street of the streets, and the sun and soul of the station
and so too too the world of a sound and stranger and to the world
to the sun a

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

Вывод

В этой статье мы узнали о RNN и LSTM. Мы также создали нашу модель стихотворения с помощью PyTorch. Надеюсь, эта статья оказалась для вас полезной. Если у вас есть какие-либо вопросы или предложения, не стесняйтесь размещать их в разделе комментариев ниже или свяжитесь со мной по адресу [email protected], я буду очень рад вам помочь.