RNN con PyTorch: red neuronal recurrente desde cero

RNN vs arquitectura feedforward

Personalmente me resulta más fácil entender las RNN cuando las comparo con redes feedforward, porque es un concepto conocido y solo voy sumando ideas nuevas; por eso las compararé a menudo.

A diferencia de las redes feedforward, el mecanismo de las RNN es un poco más complejo. Dentro de una sola capa de red neuronal recurrente hay 3 matrices de pesos, además de 2 tensores de entrada y 2 tensores de salida.

Suele decirse que las RNN son feedforward con estado interno; pero con un diagrama sencillo se ve que no es tan simple. Los componentes son bastante más complejos en una red recurrente; no te preocupes, intentaré explicarte cómo funciona y, con el código, podrás entenderlo.

Arquitectura de una capa RNN

Las redes recurrentes introducen el concepto de estado oculto: básicamente otra entrada que depende de las salidas anteriores de la capa. Si depende de salidas anteriores, ¿de dónde sale en la primera pasada? Fácil: inícialo con ceros.

Las RNN se alimentan distinto a las feedforward. Como trabajamos con secuencias, el orden en que entran los datos importa; por eso en cada paso alimentamos la red con un solo elemento de la secuencia. Por ejemplo: si son precios de acciones, entramos el precio de cada día; si es texto, entramos una letra o palabra cada vez.

Entramos paso a paso porque en cada iteración hay que calcular el estado oculto; ese estado guarda información previa para que la siguiente entrada use datos de pasadas anteriores mediante la suma de las matrices.

Entradas

• Tensor de entrada: Debe ser un solo paso de la secuencia. Si tu secuencia tiene 100 caracteres, la entrada será un solo carácter.
• Tensor de estado oculto: Es el estado oculto. En la primera pasada de cada secuencia completa, este tensor se rellena con ceros. Siguiendo el ejemplo: si tienes 10 secuencias de 100 caracteres (1000 caracteres en total), para cada secuencia generarás un estado oculto inicializado en ceros.

Matrices de pesos

• Dense de entrada: Matriz densa para la entrada (como en feedforward).
• Dense oculta: Matriz densa para la entrada del estado oculto.
• Dense de salida: Matriz densa para el resultado de activation(input_dense + hidden_dense).

activation(input_dense + hidden_dense)

Salidas

• Nuevo estado oculto: Tensor de estado oculto activation(input_dense + hidden_dense). Lo usarás como entrada en la siguiente iteración de la secuencia.
• Salida: activation(output_dense). Es tu vector de predicción, análogo a la salida de una red feedforward.

Código de una capa RNN

import torch
import torch.nn as nn

class RNN(nn.Module):
    """ Bloque RNN básico. Representa una sola capa de RNN """
    def __init__(self, input_size: int, hidden_size: int, output_size: int) -> None:
        """
        input_size: Número de características del vector de entrada
        hidden_size: Número de neuronas ocultas
        output_size: Número de características del vector de salida
        """
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        
        self.i2h = nn.Linear(input_size, hidden_size, bias=False)
        self.h2h = nn.Linear(hidden_size, hidden_size)
        self.h2o = nn.Linear(hidden_size, output_size)

    def forward(self, x, hidden_state) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Devuelve la salida calculada y tanh(i2h + h2h)
        Entradas
        -------
        x: Vector de entrada
        hidden_state: Estado oculto anterior
        Salidas
        -------
        out: Salida lineal (sin activación por cómo funciona PyTorch)
        hidden_state: Nueva matriz de estado oculto
        """
        x = self.i2h(x)
        hidden_state = self.h2h(hidden_state)
        hidden_state = torch.tanh(x + hidden_state)
        out = self.h2o(hidden_state)
        return out, hidden_state

    def init_zero_hidden(self, batch_size=1) -> torch.Tensor:
        """
        Función auxiliar. Devuelve un estado oculto con el tamaño de lote indicado. Por defecto 1.
        """
        return torch.zeros(batch_size, self.hidden_size, requires_grad=False)

Entrenamiento por lotes

Alimentar una RNN por lotes suele ser mucho más rápido (fácilmente 10x), y las RNN no son la excepción. Entrenar por lotes no mejora el rendimiento del modelo en sí: si tu red no funciona con un ejemplo a la vez, tampoco con 10 o 100.

En el ejemplo la RNN se entrena con texto, un carácter a la vez; la función de entrenamiento debe dar un carácter del texto en cada paso. Hacerlo por lotes ahorra mucho tiempo, así que se pueden procesar varios lotes por época.

def train(model: RNN, data: DataLoader, epochs: int, optimizer: optim.Optimizer, loss_fn: nn.Module) -> None:
    """ Entrena el modelo durante el número de épocas indicado """
    train_losses = {}
    model.to(device)
    model.train()
    print("=> Iniciando entrenamiento")
    for epoch in range(epochs):
        epoch_losses = list()
        for X, Y in data:
            if X.shape[0] != model.batch_size:
                continue
            hidden = model.init_zero_hidden(batch_size=model.batch_size)
            X, Y, hidden = X.to(device), Y.to(device), hidden.to(device)
            
            model.zero_grad()
            loss = 0
            for c in range(X.shape[1]):
                out, hidden = model(X[:, c].reshape(X.shape[0],1), hidden)
                l = loss_fn(out, Y[:, c].long())
                loss += l
            
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 3)
            optimizer.step()
            epoch_losses.append(loss.detach().item() / X.shape[1])
        
        train_losses[epoch] = torch.tensor(epoch_losses).mean()
        print(f'=> época: {epoch + 1}, pérdida: {train_losses[epoch]}')

Programar esto a mano ayuda mucho a entender las operaciones y el flujo; además es muy satisfactorio ver cómo la RNN aprende del texto y genera nuevo texto.