Clasificación de texto - Transfer Learning

En el post anterior vimos cómo podemos entrenar una red neuronal recurrente para clasificar texto. En este tipo de tarea, nuestro modelo será capaz de asignar una etiqueta concreta entre varias a una pieza de texto determinada. Vimos, por ejemplo, que podemos saber de manera automática si una opinión de una película es positiva o negativa. En este post vamos a resolver exactamente el mismo caso, pero introduciendo una nueva técnica muy utilizada: el transfer learning.

Transfer Learning

Esta técnica nos permite entrenar redes neuronales de manera más rápida, con menores requisitos computacionales y permitiendo el entrenamiento de redes con mejores prestaciones con pequeños datasets. La idea consiste en entrenar una red neuronal en un gran dataset, con grandes recursos computacionales, y una vez entrenada utilizar el conocimiento que este modelo ya posee como punto de partida para nuestro caso particular en el proceso conocido como fine tuning.

Este proceso de fine tuning puede variar según la tarea, pero lo más común es sustituir las capas finales de la red por nuevas capas adaptadas a nuestra tarea y entrenar solo estas nuevas capas, dejando intactas las capas ya entrenadas. Sin embargo, en el caso en el que los datos usados en la nueva tarea sean muy diferentes que los usados originalmente, también es común el entrenamiento de toda la red, a partir de los pesos pre-entrenados.

Como comentábamos al principio, esta técnica es muy utilizada en la práctica. Podemos encontrar modelos pre-entrenados en diferentes librerías, que podemos descargar y empezar a utilizar directamente. El transfer learning es utilizado tanto en aplicaciones de lenguaje como tareas visuales, y lo usaremos de manera extensiva de ahora en adelante.

El dataset

Seguimos utilizando el dataset IMDB, disponible en torchtext.

import torch
import torchtext

TEXT = torchtext.data.Field(tokenize = 'spacy')
LABEL = torchtext.data.LabelField(dtype = torch.long)

train_data, test_data = torchtext.datasets.IMDB.splits(TEXT, LABEL)

len(train_data), len(test_data)

(25000, 25000)

print(vars(train_data.examples[0]))

{'text': ['How', 'I', 'got', 'into', 'it', ':', 'When', 'I', 'started', 'watching', 'this', 'series', 'on', 'Cartoon', 'Network', ',', 'I', 'have', 'to', 'say', 'that', 'I', "'ve", 'never', 'seen', 'anything', 'like', 'this', ',', 'and', 'it', 'was', 'the', 'best', '.', 'But', 'when', 'I', 'started', 'collecting', 'the', 'series', 'on', 'VHS', ',', 'and', 'years', 'later', 'on', 'DVD', 'part', 'of', 'Bandai', "'s", 'Anime', 'Legends', 'collections', '.', 'It', 'was', 'amazing', ',', 'and', 'truly', 'worth', 'watching', '.', 'It', 'had', 'a', 'lot', 'of', 'exploding', 'action', 'that', 'will', 'blow', 'you', 'out', 'of', 'your', 'seat', '.', 'And', 'of', 'course', ',', 'the', 'theme', 'songs', '"', 'Just', 'Communication",and', 'Rhythm', 'Emotions', '"', 'were', 'the', 'best.<br', '/><br', '/>Characters', ',', 'and', 'Gundams', ':', 'My', 'favorite', 'characters', 'in', 'the', 'show', 'were', ':', 'Heero', ',', 'Duo', ',', 'Relena', ',', 'Treize', ',', 'Lady', 'Und', ',', 'Noin', ',', 'and', 'Zechs', '.', 'My', 'favorite', 'Gundams', 'in', 'the', 'show', 'that', 'I', 'liked', 'the', 'most', 'are', 'the', 'Wing', 'Zero', ',', 'and', 'Epyon', ',', 'and', 'of', 'course', 'the', 'Altron', ',', 'and', 'Deathscythe', 'I', ',', 'and', 'II.<br', '/><br', '/>Meaning', 'of', 'the', 'show', ':', 'What', 'this', 'series', 'also', 'tells', 'us', 'that', 'in', 'real', 'life', ',', 'wars', 'are', 'very', 'hard', 'and', 'we', 'can', 'sometimes', 'win', ',', 'or', 'lose', '.', 'But', 'peace', 'can', 'also', 'be', 'hard', 'to', 'obtain', ',', 'and', 'I', 'do', 'believe', 'the', 'Gundam', 'pilots', 'are', 'doing', 'the', 'right', 'thing', ',', 'and', 'are', 'trying', 'to', 'obtain', 'world', 'peace.<br', '/><br', '/>But', 'however', ',', 'this', 'show', 'is', 'truly', 'the', 'best', 'of', 'the', 'best', '.', 'So', 'in', 'closing', 'to', 'this', 'review', ',', 'after', 'you', 'watch', 'this', 'show', ',', 'see', 'the', 'Movie', 'Endless', 'Waltz', '.'], 'label': 'pos'}

Embeddings pre-entrenados

El primer ejemplo de transfer learning que vamos a ver es el uso de embeddings pre-entrenados. Recuerda que un embedding es la respresentación vectorial de cada palabra en el vocabulario que utilizaremos para alimentar nuestra red recurrente. Puedes aprender más sobre embeddings en este post. En torchtext podemos descargar estos embeddings en la función build_vocab, con el parámtero vectors. En la documentación encontrarás los diferentes embeddings disponibles.

MAX_VOCAB_SIZE = 10000

TEXT.build_vocab(train_data, 
                 max_size = MAX_VOCAB_SIZE, 
                 vectors = "glove.6B.100d", # embeddings pre-entrenados
                 unk_init = torch.Tensor.normal_)

LABEL.build_vocab(train_data)

len(TEXT.vocab), len(LABEL.vocab)

(10002, 2)

Y, de la misma manera que hicimos en el post anterior, definimos nuestros dataloaders con la clase torchtext.data.BucketIterator.

device = "cuda" if torch.cuda.is_available() else "cpu"

dataloader = {
    'train': torchtext.data.BucketIterator(train_data, batch_size=64, shuffle=True, sort_within_batch=True, device=device),
    'test': torchtext.data.BucketIterator(test_data, batch_size=64, device=device)
}

El modelo

Usaremos exactamente el mismo modelo que ya vimos en el post anterior. Este modelo está compuesto, principalmente, por la capa embedding, que en este caso sustituiremos por los vectores descargados anteriormente, y las capas recurrente y lineal, que entrenaremos desde cero.

class RNN(torch.nn.Module):
    def __init__(self, input_dim, embedding_dim=128, hidden_dim=128, output_dim=2, num_layers=2, dropout=0.2, bidirectional=False):
        super().__init__()
        self.embedding = torch.nn.Embedding(input_dim, embedding_dim)
        self.rnn = torch.nn.GRU(
            input_size=embedding_dim, 
            hidden_size=hidden_dim, 
            num_layers=num_layers, 
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=bidirectional
        )
        self.fc = torch.nn.Linear(2*hidden_dim if bidirectional else hidden_dim, output_dim)
        
    def forward(self, text):
        # no entrenamos los embeddings
        with torch.no_grad():
            #text = [sent len, batch size]        
            embedded = self.embedding(text)        
        #embedded = [sent len, batch size, emb dim]        
        output, hidden = self.rnn(embedded)        
        #output = [sent len, batch size, hid dim]
        y = self.fc(output[-1,:,:].squeeze(0))     
        return y

Una vez definido el modelo, sustituimos los tensores en la capa embedding por los vectores pre-entrenados descargados anteriormente.

model = RNN(input_dim=len(TEXT.vocab), bidirectional=True, embedding_dim=100)

pretrained_embeddings = TEXT.vocab.vectors
model.embedding.weight.data.copy_(pretrained_embeddings)
# ponemos a cero los pesos correspondientes a los tokens <unk> y <pad>
model.embedding.weight.data[TEXT.vocab.stoi[TEXT.unk_token]] = torch.zeros(100)
model.embedding.weight.data[TEXT.vocab.stoi[TEXT.pad_token]] = torch.zeros(100)

outputs = model(torch.randint(0, len(TEXT.vocab), (100, 64)))
outputs.shape

torch.Size([64, 2])

Entrenamiento

Para entrenar nuestra red usamos el bucle estándar que ya usamos en posts anteriores.

from tqdm import tqdm
import numpy as np

def fit(model, dataloader, epochs=5):
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = torch.nn.CrossEntropyLoss()
    for epoch in range(1, epochs+1):
        model.train()
        train_loss, train_acc = [], []
        bar = tqdm(dataloader['train'])
        for batch in bar:
            X, y = batch
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            y_hat = model(X)
            loss = criterion(y_hat, y)
            loss.backward()
            optimizer.step()
            train_loss.append(loss.item())
            acc = (y == torch.argmax(y_hat, axis=1)).sum().item() / len(y)
            train_acc.append(acc)
            bar.set_description(f"loss {np.mean(train_loss):.5f} acc {np.mean(train_acc):.5f}")
        bar = tqdm(dataloader['test'])
        val_loss, val_acc = [], []
        model.eval()
        with torch.no_grad():
            for batch in bar:
                X, y = batch
                X, y = X.to(device), y.to(device)
                y_hat = model(X)
                loss = criterion(y_hat, y)
                val_loss.append(loss.item())
                acc = (y == torch.argmax(y_hat, axis=1)).sum().item() / len(y)
                val_acc.append(acc)
                bar.set_description(f"val_loss {np.mean(val_loss):.5f} val_acc {np.mean(val_acc):.5f}")
        print(f"Epoch {epoch}/{epochs} loss {np.mean(train_loss):.5f} val_loss {np.mean(val_loss):.5f} acc {np.mean(train_acc):.5f} val_acc {np.mean(val_acc):.5f}")

fit(model, dataloader)

  0%|          | 0/391 [00:00<?, ?it/s]
loss 0.60385 acc 0.65508: 100%|██████████| 391/391 [00:29<00:00, 13.47it/s]
val_loss 1.20899 val_acc 0.50868: 100%|██████████| 391/391 [00:28<00:00, 13.63it/s]
loss 0.41505 acc 0.85156:   0%|          | 1/391 [00:00<00:42,  9.20it/s]

Epoch 1/5 loss 0.60385 val_loss 1.20899 acc 0.65508 val_acc 0.50868


loss 0.38606 acc 0.82816: 100%|██████████| 391/391 [00:28<00:00, 13.50it/s]
val_loss 0.36242 val_acc 0.83732: 100%|██████████| 391/391 [00:28<00:00, 13.79it/s]
loss 0.29536 acc 0.88281:   0%|          | 1/391 [00:00<00:46,  8.34it/s]

Epoch 2/5 loss 0.38606 val_loss 0.36242 acc 0.82816 val_acc 0.83732


loss 0.31376 acc 0.86516: 100%|██████████| 391/391 [00:29<00:00, 13.48it/s]
val_loss 0.37173 val_acc 0.83473: 100%|██████████| 391/391 [00:29<00:00, 13.44it/s]
  0%|          | 0/391 [00:00<?, ?it/s]

Epoch 3/5 loss 0.31376 val_loss 0.37173 acc 0.86516 val_acc 0.83473


loss 0.27493 acc 0.88288: 100%|██████████| 391/391 [00:29<00:00, 13.42it/s]
val_loss 0.32504 val_acc 0.86353: 100%|██████████| 391/391 [00:28<00:00, 13.71it/s]
loss 0.23483 acc 0.89062:   1%|          | 2/391 [00:00<00:31, 12.42it/s]

Epoch 4/5 loss 0.27493 val_loss 0.32504 acc 0.88288 val_acc 0.86353


loss 0.23740 acc 0.90324: 100%|██████████| 391/391 [00:28<00:00, 13.50it/s]
val_loss 0.31196 val_acc 0.86348: 100%|██████████| 391/391 [00:28<00:00, 13.58it/s]

Epoch 5/5 loss 0.23740 val_loss 0.31196 acc 0.90324 val_acc 0.86348

Generando predicciones

Una vez nuestro modelo ha sido entrenado, podemos generar predicciones exactamente igual que hicimos en el post anterior.

import spacy
nlp = spacy.load('en')

def predict(model, X):
    model.eval() 
    with torch.no_grad():
        X = torch.tensor(X).to(device)
        pred = model(X)
        return pred

sentences = ["this film is terrible", "this film is great", "this film is good", "a waste of time"]
tokenized = [[tok.text for tok in nlp.tokenizer(sentence)] for sentence in sentences]
indexed = [[TEXT.vocab.stoi[_t] for _t in t] for t in tokenized]
tensor = torch.tensor(indexed).permute(1,0)
predictions = torch.argmax(predict(model, tensor), axis=1)
predictions

tensor([0, 1, 1, 0], device='cuda:0')

Transformers

Como has visto en el ejemplo anterior, utilizar unos embeddings pre-entrenados puede darnos mucho mejores resultados que entrenarlos desde cero, ya que la representación de nuestras palabras será mucho mejor desde el principio. Siguiendo en esta línea, podemos sustituir nuestra capa embedding por otro modelo que nos aportará todavía mejores resultados, un transformer.

Estos modelos aparecieron alrededor de 2017, y fueron presentados en el famoso artículo Attention is All You Need. Desde su aparación, estos modelos están batiendo todos los benchmarks en las diferentes tareas de procesado de lenguaje, y son utilizados como base de cualquier modelo competente a día de hoy. De momento, no entraremos en detalles en la definición de esta arquitectura (lo dejamos para un futuro post, ya que hay mucha tela que cortar) pero vamos a ver como utilizar un transformer para hacer transfer learning y obtener muy buenos resultados de manera rápida.

Una librería muy utilizada para trabajar con estos modelos es la librería transformers de huggingface.

!pip install transformers

En primer lugar, tendremos que utilizar el mismo tokenizer utilizado para entrenar el modelo original. En este caso usaremos la red conocida como BERT.

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

tokens = tokenizer.tokenize('Hello WORLD how ARE yoU?')
tokens

['hello', 'world', 'how', 'are', 'you', '?']

indexes = tokenizer.convert_tokens_to_ids(tokens)
indexes

[7592, 2088, 2129, 2024, 2017, 1029]

A diferencia de las redes neuronales recurrentes, los transformers trabajan con longitudes de secuencia fijas (no son modelos recurrentes). Es por este motivo que tenemos que asegurarnos que ninguna frase en el dataset tiene mayor longitud que la máxima permitida por BERT, que es de 512 palabras.

max_input_length = tokenizer.max_model_input_sizes['bert-base-uncased']

def tokenize_and_cut(sentence):
    tokens = tokenizer.tokenize(sentence) 
    tokens = tokens[:max_input_length-2]
    return tokens

torchtext nos da la libertad de definir nuestros propios tokenizers, y podemos incluirlos de la siguiente manera.

TEXT = torchtext.data.Field(batch_first = True,
                  use_vocab = False,
                  tokenize = tokenize_and_cut,
                  preprocessing = tokenizer.convert_tokens_to_ids,
                  init_token = tokenizer.cls_token_id,
                  eos_token = tokenizer.sep_token_id,
                  pad_token = tokenizer.pad_token_id,
                  unk_token = tokenizer.unk_token_id)

LABEL = torchtext.data.LabelField(dtype = torch.long)

train_data, test_data = torchtext.datasets.IMDB.splits(TEXT, LABEL)

LABEL.build_vocab(train_data)

dataloader = {
    'train': torchtext.data.BucketIterator(train_data, batch_size=64, shuffle=True, sort_within_batch=True, device=device),
    'test': torchtext.data.BucketIterator(test_data, batch_size=64, device=device)
}

Una vez tenemos los datos preparados con el nuevo tokenizer, necesitamos definir nuestro nuevo modelo. En este caso, BERT se encargará de actuar como nuestra capa embedding, proveyendo de la mejor representación posible de nuestro texto para que las siguientes capas puedan clasificarlo.

from transformers import BertModel

class BERT(torch.nn.Module):
    def __init__(self, hidden_dim=256, output_dim=2, n_layers=2, bidirectional=True, dropout=0.2):
        super().__init__()        
        self.bert = BertModel.from_pretrained('bert-base-uncased')        
        
        # freeze BERT
        for name, param in self.bert.named_parameters():                
            if name.startswith('bert'):
                param.requires_grad = False

        embedding_dim = self.bert.config.to_dict()['hidden_size']
        self.rnn = torch.nn.GRU(embedding_dim,
                          hidden_dim,
                          num_layers = n_layers,
                          bidirectional = bidirectional,
                          batch_first = True,
                          dropout = 0 if n_layers < 2 else dropout)
        
        self.fc = torch.nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, output_dim)
        
    def forward(self, text):                       
        with torch.no_grad():
            embedded = self.bert(text)[0]
        output, hidden = self.rnn(embedded)        
        y = self.fc(output[:,-1,:].squeeze(1))     
        return y

model = BERT()
fit(model, dataloader, epochs=3)

  0%|          | 0/391 [00:00<?, ?it/s]
loss 0.35525 acc 0.83775: 100%|██████████| 391/391 [04:41<00:00,  1.39it/s]
val_loss 0.24976 val_acc 0.90348: 100%|██████████| 391/391 [07:33<00:00,  1.16s/it]
  0%|          | 0/391 [00:00<?, ?it/s]

Epoch 1/3 loss 0.35525 val_loss 0.24976 acc 0.83775 val_acc 0.90348


loss 0.22018 acc 0.91375: 100%|██████████| 391/391 [04:41<00:00,  1.39it/s]
val_loss 0.29960 val_acc 0.86585: 100%|██████████| 391/391 [07:33<00:00,  1.16s/it]
  0%|          | 0/391 [00:00<?, ?it/s]

Epoch 2/3 loss 0.22018 val_loss 0.29960 acc 0.91375 val_acc 0.86585


loss 0.18880 acc 0.92842: 100%|██████████| 391/391 [04:41<00:00,  1.39it/s]
val_loss 0.44130 val_acc 0.77996: 100%|██████████| 391/391 [07:33<00:00,  1.16s/it]
  0%|          | 0/391 [00:00<?, ?it/s]

Epoch 3/3 loss 0.18880 val_loss 0.44130 acc 0.92842 val_acc 0.77996

Puedes ver que en la primera epoch ya tenemos un modelo mejor que cualquiera de los obtenidos anteriormente cuando entrenamos desde cero. Y finalmente, podemos generar las predicciones de la siguiente manera

def predict(sentence):
    tokenized = [tok[:max_input_length-2] for tok in tokenizer.tokenize(sentence)]
    indexed = [tokenizer.cls_token_id] + tokenizer.convert_tokens_to_ids(tokenized) + [tokenizer.sep_token_id]
    tensor = torch.tensor([indexed]).to(device)
    model.eval()
    return torch.argmax(model(tensor), axis=1)

sentences = ["Best film ever !", "this movie is terrible"]
preds = [predict(s) for s in sentences]
preds

[tensor([1], device='cuda:0'), tensor([0], device='cuda:0')]

Resumen

En este post hemos visto como podemos obtener mejores modelos de clasificación de texto si utilizamos el transfer learning. Con esta técnica, sustituiremos las primeras capas de nuestro modelo por otros modelos ya entrenados en otros datasets. Dadas las condiciones adecuadas, esta técnica nos va a permitir entrenar modelos muy buenos con pocos datos y en poco tiempo. En nuestro caso, hemos conseguido nuestro mejor clasificador utilizando el modelo BERT como capa embedding y entrenando nuestro modelo recurrente encima, el cual solo hemos necesitado entrenar por una epoch.