[pytorch] 순환신경망 (Recurrent Neural Network) 예제 코드

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

n_hidden = 35
lr = 0.01
epochs = 1000

string = "hello pytorch. how long can a rnn cell remember? show me your limit!"
chars =  "abcdefghijklmnopqrstuvwxyz ?!.,:;01"
char_list = [i for i in chars]
n_letters = len(char_list)

# one-hot encoding
def string_to_onehot(string):
  start = np.zeros(shape=n_letters, dtype=int)
  end = np.zeros(shape=n_letters, dtype=int)
  start[-2] = 1 # start 토큰은 [0,0, ..., 1, 0]
  end[-1] =1 # end 토큰은 [0, 0, ..., 0, 1]
  for i in string:
    idx = char_list.index(i)
    zero = np.zeros(shape=n_letters, dtype=int)
    zero[idx] = 1
    start = np.vstack([start,zero]) # vertical하게 벡터 결합
  output = np.vstack([start, end])
  return output

# one-hot to word
def onehot_to_word(onehot_1):
  onehot = torch.Tensor.numpy(onehot_1)
  return char_list[onehot.argmax()]


# 모델 재사용성을 위해 클래스로 설계
class RNN(nn.Module):
  def __init__(self, input_size, hidden_size, output_size):
    super(RNN, self).__init__()

    self.input_size = input_size
    self.hidden_size = hidden_size
    self.output_size = output_size

    self.i2h = nn.Linear(input_size, hidden_size)
    self.h2h = nn.Linear(hidden_size, hidden_size)
    self.i2h = nn.Linear(hidden_size, output_size)
    self.act_fn = nn.Tanh()

  def forward(self, input, hidden):
    hidden = self.act_fn(self.i2h(input)+self.h2h(hidden))
    output = self.i2h(hidden)
    return output, hidden

  def init_hidden(self):
    return torch.zeros(1, self.hidden_size)


# Model
rnn = RNN(n_letters, n_hidden, n_letters)

# loss, optim
loss_func = nn.MSELoss()
optimizer = torch.optim.Adam(rnn.parameters(), lr=lr)


# Training
one_hot = torch.from_numpy(string_to_onehot(string)).type_as(torch.FloatTensor())

for i in range(epochs):
    optimizer.zero_grad()
    # t=0 이므로 hidden 초기화
    hidden = rnn.init_hidden()
    
    total_loss = 0
    for j in range(one_hot.size()[0]-1):
        # 입력은 앞에 글자 
        # pyotrch 에서 p y t o r c
        input_ = one_hot[j:j+1,:]
        # 목표값은 뒤에 글자
        # pytorch 에서 y t o r c h
        target = one_hot[j+1]
        output, hidden = rnn.forward(input_, hidden)
        
        loss = loss_func(output.view(-1),target.view(-1))
        total_loss += loss

    total_loss.backward()
    optimizer.step()

    if i % 10 == 0:
        print(total_loss)


# Test 
start = torch.zeros(1,n_letters)
start[:,-2] = 1

with torch.no_grad():
    hidden = rnn.init_hidden()
    input_ = start
    output_string = ""

    for i in range(len(string)):
        output, hidden = rnn.forward(input_, hidden)
        output_string += onehot_to_word(output.data)
        input_ = output

print(output_string)