Distilling BERT to simpler models: CNN and Linear¶

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import altair as alt
from pathlib import Path
from joblib import Memory
from sklearn.metrics import f1_score
from tqdm import tqdm
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig, DistilBertForSequenceClassification

ROOT = Path("/mnt/gdrive/My Drive/")

if not ROOT.exists():
    from google.colab import drive
    drive.mount("/mnt/gdrive")

assert ROOT.exists()
DATA = ROOT / "data"
MODELS = ROOT / "models"

CACHE = ROOT / "cache/distillation"

if not CACHE.exists():
    CACHE.mkdir(parents=True)

memory = Memory(CACHE, verbose=False)

market = "uk"

reviews = (pd.read_csv(DATA / "amazon" / f"amazon_reviews_multilingual_{market.upper()}_v1_00.tsv.gz",
                       sep="\t",
                       usecols=["review_id", "star_rating", "review_headline", "review_body"],
                       dtype={"review_id": "string",
                              "star_rating": "Int32",
                              "review_headline": "string",
                              "review_body": "string"})
            .dropna())

try:
    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-multilingual-cased")
except NameError:
    tokenizer = tokenizer

@memory.cache(ignore=["teacher"])  # cache but don't bother serializing the teacher - doesn't change and is slow to pickle
def dataframe_to_dataset(df, teacher):
    max_len = 128
    features = tokenizer.batch_encode_plus(df.review_body,
                                           max_length=max_len,
                                           pad_to_max_length=True,
                                           return_attention_masks=True,
                                           return_token_type_ids=True,
                                           return_tensors="pt")

    pre_dataset = TensorDataset(features["input_ids"],
                                features["attention_mask"],
                                features["token_type_ids"])
    
    teacher.to(device)
    teacher.eval()
    teacher_predictions = []
    for batch in tqdm(DataLoader(pre_dataset, batch_size=32, shuffle=False)):
        batch = tuple([b.to(device) for b in batch])
        inputs = {"input_ids": batch[0], "attention_mask": batch[1]}
        if teacher.base_model_prefix == "bert":
            inputs["token_type_ids"] = batch[2]
        with torch.no_grad():
            outputs = teacher(**inputs)
        teacher_predictions.append(outputs[0].to(torch.device("cpu")))  # put back on CPU

    dataset = TensorDataset(features["input_ids"],
                            features["attention_mask"],
                            features["token_type_ids"],
                            torch.tensor(df.label.astype("int").to_numpy(), dtype=torch.long),
                            torch.cat(teacher_predictions, axis=0))
    return dataset

N_EPOCHS = 5
EMBEDDING_DIM = 50
LR = 0.5
BATCH_SIZE = 32
N_LABELS = 5  # num review ratings

padding_idx = tokenizer.vocab["[PAD]"]
n_vocab = len(tokenizer.vocab)

try:
    config = config
    teacher = teacher
except NameError:
    config = AutoConfig.from_pretrained("distilbert-base-multilingual-cased")
    config.num_labels = N_LABELS
    teacher = DistilBertForSequenceClassification(config)
    teacher.load_state_dict(torch.load(MODELS / "distilbert_uk_50000.bin", map_location=device))

class CNN(nn.Module):

    def __init__(self,
                 n_vocab,
                 n_labels,
                 embedding_dim=50,
                 n_filters=100,
                 filter_sizes=[3, 4, 5],
                 dropout=0.5,
                 special_chars=[],
                 pretrained_embeddings=None):  # TODO make number of conv layers configurable
        super(CNN, self).__init__()
        self.n_vocab = n_vocab
        self.n_labels = n_labels
        self.embedding_dim = embedding_dim
        self.n_filters = n_filters
        self.filter_sizes = filter_sizes
        self.dropout_p = dropout
        self.width = len(filter_sizes) * n_filters

        if pretrained_embeddings is not None:
            assert n_vocab == pretrained_embeddings.shape[0]
            assert embedding_dim == pretrained_embeddings.shape[1]
            self.embedding = nn.Embedding.from_pretrained(pretrained_embeddings)
        else:
            self.embedding = nn.Embedding(n_vocab, embedding_dim)
        
        self.conv0 = nn.Conv2d(in_channels=1,
                               out_channels=n_filters,
                               kernel_size=(filter_sizes[0], embedding_dim))
        self.conv1 = nn.Conv2d(in_channels=1,
                               out_channels=n_filters,
                               kernel_size=(filter_sizes[1], embedding_dim))
        self.conv2 = nn.Conv2d(in_channels=1,
                               out_channels=n_filters,
                               kernel_size=(filter_sizes[2], embedding_dim))
        self.dropout = nn.Dropout(dropout)

        self.fc = nn.Linear(in_features=self.width, out_features=n_labels)

        for special in special_chars:
            self.embedding.weight.data[special] = torch.zeros(embedding_dim)

    def forward(self, input_ids, **kwargs):
        """Only input ids are required - kwargs are for API compat with BERT."""
        X = self.embedding(input_ids)
        X = X.unsqueeze(1)  # add single channel as dim 1
        X0 = F.relu(self.conv0(X).squeeze(3))
        X1 = F.relu(self.conv1(X).squeeze(3))
        X2 = F.relu(self.conv2(X).squeeze(3))
        X0 = F.max_pool1d(X0, X0.shape[2]).squeeze(2)
        X1 = F.max_pool1d(X1, X1.shape[2]).squeeze(2)
        X2 = F.max_pool1d(X2, X2.shape[2]).squeeze(2)
        X = torch.cat([X0, X1, X2], dim=1)
        X = self.dropout(X)
        X = self.fc(X)
        return X

class LinearModel(nn.Module):
    
    def __init__(self, n_vocab, n_labels, embedding_dim, padding_idx):
        super(LinearModel, self).__init__()
        self.embeddings = nn.Embedding(n_vocab, embedding_dim, padding_idx=padding_idx)
        self.output = nn.Linear(embedding_dim, n_labels)
        with torch.no_grad():
            # FastText initializes embeddings with uniform distribution vs normal in PyTorch
            self.embeddings.weight.uniform_(to=1.0 / embedding_dim)
            self.embeddings.weight[padding_idx] = 0  # but FT doesn't have a padding token
            # FastText initializes output with zeros vs some random dist in PyTorch
            self.output.weight.zero_()

    def forward(self, input_ids, **kwargs):
        """Only input ids are required - kwargs are for API compat with BERT."""
        X = self.embeddings(input_ids)
        X = X.mean(dim=1)
        X = self.output(X)
        return X

def train_epoch(train_iter, model, optim, epoch_num, distil=False):
    train_loss = 0
    train_acc = 0
    y_true = []
    y_pred = []
    
    model.to(device)
    model.train()
    
    if distil:
        cost = nn.MSELoss()
    else:
        cost = nn.CrossEntropyLoss()

    for batch in tqdm(train_iter, total=len(train_iter), desc=f"Batch progress for epoch {epoch_num}"):
        
        batch = tuple([t.to(device) for t in batch])
        inputs = {"input_ids": batch[0],
                  "attention_mask": batch[1]}
        labels = batch[3]

        optim.zero_grad()
        output = model(**inputs)

        if distil:
            target = batch[4]
        else:
            target = labels

        batch_loss = cost(output, target)

        # Had some trouble with linear distilled model dying in training
        # but since starting to debug it the issue hasn't reoccurred.
        # Gradient clipping might help. 
        if torch.isnan(batch_loss):
            print("NAN batch loss!", epoch_num, batch_loss, output, target)

        train_loss += batch_loss.item()

        batch_acc = (output.argmax(1) == labels).sum().item()
        train_acc += batch_acc
        y_true.extend(labels.tolist())
        y_pred.extend(output.argmax(1).tolist())

        batch_loss.backward()
        optim.step()

    return train_loss / len(train_iter), train_acc / len(train_iter.dataset), f1_score(y_true, y_pred, average="macro")  # classes are already balanced

def train_loop(model, optim, train_loader, test_loader, n_epochs=5, sched=None, distil=False):
    training_results = {"epoch": list(range(n_epochs)),
                        "train_loss": [],
                        "train_acc": [],
                        "train_f1_macro": [],
                        "test_loss": [],
                        "test_acc": [],
                        "test_f1_macro": []}

    model.to(device)

    try:
        for i in range(n_epochs):
            train_loss, train_acc, train_f1 = train_epoch(train_loader, model, optim, epoch_num=i, distil=distil)
            if sched is not None:
                sched.step()
            test_loss, test_acc, test_f1 = validate(test_loader, model)
            training_results["train_loss"].append(train_loss)
            training_results["train_acc"].append(train_acc)
            training_results["train_f1_macro"].append(train_f1)
            training_results["test_loss"].append(test_loss)
            training_results["test_acc"].append(test_acc)
            training_results["test_f1_macro"].append(test_f1)
    except KeyboardInterrupt:
        pass

    return pd.DataFrame(training_results)

def validate(test_iter, model):
    test_acc = 0 
    test_loss = 0
    y_true = []
    y_pred = []

    cost = nn.CrossEntropyLoss()

    model.to(device)
    model.eval()

    for batch in tqdm(test_iter, desc="Validating"):
        
        batch = tuple([t.to(device) for t in batch])
        inputs = {"input_ids": batch[0],
                  "attention_mask": batch[1],
                  "token_type_ids": batch[2]}
        labels = batch[3]

        with torch.no_grad():
            output = model(**inputs)
            
            batch_loss = cost(output, labels)
            test_loss += batch_loss.item()
                    
            batch_acc = (output.argmax(1) == labels).sum().item() 
            test_acc += batch_acc
            y_true.extend(labels.tolist())
            y_pred.extend(output.argmax(1).tolist())

    return test_loss / len(test_iter), test_acc / len(test_iter.dataset), f1_score(y_true, y_pred, average="macro")  # classes are balanced

linear_model = LinearModel(n_vocab, N_LABELS, embedding_dim=EMBEDDING_DIM, padding_idx=padding_idx)

linear_model_dist = LinearModel(n_vocab, N_LABELS, embedding_dim=EMBEDDING_DIM, padding_idx=padding_idx)

cnn = CNN(n_vocab, N_LABELS, embedding_dim=EMBEDDING_DIM, special_chars=[padding_idx])

cnn_dist = CNN(n_vocab, N_LABELS, embedding_dim=EMBEDDING_DIM, special_chars=[padding_idx])

train_loader = DataLoader(dataframe_to_dataset(train_df, teacher), batch_size=BATCH_SIZE, shuffle=False)  # optimal training for the cnn

training_results = pd.concat([linear_model_train_results,
                              linear_model_dist_train_results,
                              cnn_train_results,
                              cnn_dist_train_results])