pytorch on Marcin Zabłocki blog

Robust Subspace Recovery Layer implementation in PyTorch

Marcin Zabłocki — Mon, 19 Apr 2021 00:00:00 +0200

Anomaly detection (or outlier detection) is one of the many tasks that keep Data Scientists awake at night when then need to satisfy the business needs. No matter whether the data in well structured (like user events) or unstructured (like text or images), there are plenty of methods that allow to tackle this problem and it’s diffictult to choose one. Sometimes, it’s even difficult to properly define, what an anomaly is - but that’s not the part of this post’s main focus. Some of the most popular methods for anomaly detection are:

statistical methods - like three-sigma rule and regression models
classical machine learning methods - like one-class SVM, Isolation Forests or clustering-based methods (e.g. kNN)
deep learning methods - like autoencoders and deep forecasting models (e.g. LSTMs for time series)

In this post I will focus on a single method called Robust Subspace Recovery Layer for Unsupervised Anomaly Detection (https://arxiv.org/abs/1904.00152) which is an extension of the autoencoder-based unsupervised anomnaly detection method for data of any domain.

TL;DR

Implementation of RSR Layer and RSR Autoencoder (RSRAE) in PyTorch. Link to Jupyter Notebook / Google Colab at the bottom.

Note: This post and my implementation does not reproduce RSR Layer paper. I only explain the general idea and provide minimal working code as a form of proof-of-concept.

Idea behind Robust Space Recovery Layer

Autoencoders were used for the anomaly detection much earlier than this (RSR) paper appeared. To use autoencoder in a setting of anomaly detection, one treats “normal” examples (inliers) as the ones that have low reconstruction error and other which have high reconstruction error are treated as outliers. Although the idea is fairly easy to understand and implement, the problem lies in the fact that in order for autoencoder to work properly in this setting (i.e. to detect the outliers) the training data should be clean (or have really small fraction of the outliers). In the real world it’s highly unlikely for the input data to be free of any anomalies and even if Data Scientists could label and filter the data (given they know what an anomaly in the given dataset is) it simply won’t scale.

The Robust Subspace Recovery Layer paper tries to adress this issue by incorporating the RSR Layer into autoencoders. It’s clearly stated that this method does not compete with semi-supervised or supervised anomaly detection setting, but shows its strenght in the unsupervised setting.

What is the RSR Layer?

The RSR Layer itself is just an additional layer in the nerual network with specific properties. It maps one representation (embedding obtained from the encoder) into another low-dimensional representation that is outlier-robust.

Let’s say that D (upper case D) is the dimension of the embedding from the encoder. The assumption in the paper is that the “normal” data lies within d-dimensional (lower case d) manifold (“subspace”) of the original embedding, which means that d < D.

Please refer to the paper to obtain more details and the full explanation and analysis. At this point you should know enough to proceed to the implementation section.

How to implement the RSR Layer in PyTorch?

As RSR Layer is just a linear mapping, there are only a few things needed:

The dimension D from the encoder’s output.
The dimension d which is a direct property of the RSR Layer (and probably requires a hyperparamer search to set it properly).
Proper initialization of the linear mapping weights (the optimization goal of the RSR Autoencoder enforces that, so it’s better to start with a set of weights that already fulfill this constraint).

The implementation of the layer itself is straightforward.

python

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

python

class RSRLayer(nn.Module):
    def __init__(self, d:int, D: int):
        super().__init__()
        self.d = d
        self.D = D
        self.A = nn.Parameter(torch.nn.init.orthogonal_(torch.empty(d, D)))

    def forward(self, z):
        # z is the output from the encoder
        z_hat = self.A @ z.view(z.size(0), self.D, 1)
        return z_hat.squeeze(2)

The next piece to obtain RSR Autoencoder in PyTorch is to implement RSR Loss as per paper’s equation (4): The first term enforces the RSR Layer projection to be robust and the second term enforces the projection to be orthogonal. Refer to the paper for the motivation behind this loss as it it out of scope for this blog post.

Nearly op-for-op (math-to-PyTorch) implementation of this loss looks like this:

python

class RSRLoss(nn.Module):
    def __init__(self, lambda1, lambda2, d, D):
        super().__init__()
        self.lambda1 = lambda1
        self.lambda2 = lambda2
        self.d = d
        self.D = D
        self.register_buffer(
            "Id", torch.eye(d)
        )

    def forward(self, z, A):
        z_hat = A @ z.view(z.size(0), self.D, 1)
        AtAz = (A.T @ z_hat).squeeze(2)
        term1 = torch.sum(
            torch.norm(z - AtAz, p=2)
        )

        term2 = torch.norm(
            A @ A.T - self.Id, p=2
        ) ** 2

        return self.lambda1 * term1 + self.lambda2 * term2

Lastly, the standard L1/L2 reconstruction loss is replaced by the L_2,p-norm based loss (again, the paper relies on “more robustness”, wiki: https://en.wikipedia.org/wiki/Matrix_norm#L2,1_and_Lp,q_norms). In practice, the paper uses L_2,1 version.

python

class L2p_Loss(nn.Module):
    def __init__(self, p=1.0):
        super().__init__()
        self.p = p
  
    def forward(self, y_hat, y):
        return torch.sum(
            torch.pow(
                torch.norm(y - y_hat, p=2), self.p
            )
        )

RSR Auto Encoder in PyTorch

In order to see the RSR Layer in practice, the RSR Autoencoder needs to be implemented.

Important
The following RSR Autoencoder implementation is only a minimal, working proof of concept. As I’ve stated earlier, the intent of this blog post is to implement the RSR Layers and play around with them, not reach the SOTA or reproduce the paper’s results. Feel free to use my code and plug-in your encoder and decoder networks.

python

class RSRAutoEncoder(nn.Module):
    def __init__(self, input_dim, d, D):
        super().__init__()
        # Put your encoder network here, remember about the output D-dimension
        self.encoder = nn.Sequential(
          nn.Linear(input_dim, input_dim // 2),
          nn.LeakyReLU(),
          nn.Linear(input_dim // 2, input_dim // 4),
          nn.LeakyReLU(),
          nn.Linear(input_dim // 4, D)
        )

        self.rsr = RSRLayer(d, D)

        # Put your decoder network here, rembember about the input d-dimension
        self.decoder = nn.Sequential(
          nn.Linear(d, D),
          nn.LeakyReLU(),
          nn.Linear(D, input_dim // 2),
          nn.LeakyReLU(),
          nn.Linear(input_dim // 2, input_dim)
        )
    
    def forward(self, x):
        enc = self.encoder(x) # obtain the embedding from the encoder
        latent = self.rsr(enc) # RSR manifold
        dec = self.decoder(latent) # obtain the representation in the input space
        return enc, dec, latent, self.rsr.A

Training the autoencoder with RSR Layer in PyTorch (& PyTorch Lightning)

As usual, I use PyTorch Lightning to skip the training boilerplate.

python

import pytorch_lightning as pl
pl.seed_everything(666)

Corrupted MNIST dataset (with outliers)

For this proof-of-concept I will be using a corrupted version of MNIST dataset, that contains single class as “normal” data and a few other, that will be treated as outliers.

python

from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
# https://github.com/pytorch/vision/issues/1938#issuecomment-594623431
from six.moves import urllib
opener = urllib.request.build_opener()
opener.addheaders = [('User-agent', 'Mozilla/5.0')]
urllib.request.install_opener(opener)

mnist = MNIST(".", download=True, transform=ToTensor())

python

class RSRDs(torch.utils.data.Dataset):
    def __init__(self, target_class, other_classes, n_examples_per_other):
        super().__init__()
        self.mnist = MNIST(".", download=True, transform=ToTensor())
        self.target_indices = (self.mnist.targets == target_class).nonzero().flatten()
        
        other = []
        for other_class in other_classes:
            other.extend(
                (self.mnist.targets == other_class).nonzero().flatten()[:n_examples_per_other]
            )
        self.other_indices = torch.tensor(other)
        self.all_indices = torch.cat([self.other_indices, self.target_indices])
        print(f"Targets: {self.target_indices.size(0)}")
        print(f"Others : {self.other_indices.size(0)}")
    
    def __getitem__(self, idx):
        actual_idx = self.all_indices[idx].item()
        return self.mnist[actual_idx]
    
    def __len__(self):
        return self.all_indices.size(0)

python

ds = RSRDs(target_class=4, other_classes=(0, 1, 2, 8), n_examples_per_other=100)

Targets: 5842
Others : 400

python

ds[0][0].shape, ds[0][1]

(torch.Size([1, 28, 28]), 0)

RSRAutoencoder module in PyTorch Lightning

The following module combines all of the pieces together:

RSR Autoencoder with RSR Layer inside,
L2p reconstruction loss,
RSR Loss.

python

class RSRAE(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()
        self.hparams = hparams
        self.ae = RSRAutoEncoder(
            self.hparams.input_dim, 
            self.hparams.d, 
            self.hparams.D)
        self.reconstruction_loss = L2p_Loss(p=1.0)
        self.rsr_loss = RSRLoss(self.hparams.lambda1, self.hparams.lambda2, self.hparams.d, self.hparams.D)
  
    def forward(self, x):
        return self.ae(x)
  
    def training_step(self, batch, batch_idx):
        X, _ = batch
        x = X.view(X.size(0), -1)
        enc, dec, latent, A = self.ae(x)

        rec_loss = self.reconstruction_loss(torch.sigmoid(dec), x)
        rsr_loss = self.rsr_loss(enc, A)
        loss = rec_loss + rsr_loss
        
        # log some usefull stuff
        self.log("reconstruction_loss", rec_loss.item(), on_step=True, on_epoch=False, prog_bar=True)
        self.log("rsr_loss", rsr_loss.item(), on_step=True, on_epoch=False, prog_bar=True)
        return {"loss": loss}

    def configure_optimizers(self):
        opt = torch.optim.AdamW(self.parameters(), lr=self.hparams.lr)
        # Fast.AI's best practices :)
        scheduler = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=self.hparams.lr, 
                                                        epochs=self.hparams.epochs, 
                                                        steps_per_epoch=self.hparams.steps_per_epoch)
        return [opt], [{
            "scheduler": scheduler,
            "interval": "step"
        }]

python

dl = DataLoader(ds, batch_size=64, shuffle=True, drop_last=True)

python

hparams = dict(
    d=16,
    D=128,
    input_dim=28*28,
    # Peak learning rate
    lr=0.01,
    # Configuration for the OneCycleLR scheduler
    epochs=150,
    steps_per_epoch=len(dl),
    # lambda coefficients from RSR Loss
    lambda1=1.0,
    lambda2=1.0,
)
model = RSRAE(hparams)
model

python

RSRAE(
    (ae): RSRAutoEncoder(
    (encoder): Sequential(
        (0): Linear(in_features=784, out_features=392, bias=True)
        (1): LeakyReLU(negative_slope=0.01)
        (2): Linear(in_features=392, out_features=196, bias=True)
        (3): LeakyReLU(negative_slope=0.01)
        (4): Linear(in_features=196, out_features=128, bias=True)
    )
    (rsr): RSRLayer()
    (decoder): Sequential(
        (0): Linear(in_features=16, out_features=128, bias=True)
        (1): LeakyReLU(negative_slope=0.01)
        (2): Linear(in_features=128, out_features=392, bias=True)
        (3): LeakyReLU(negative_slope=0.01)
        (4): Linear(in_features=392, out_features=784, bias=True)
    )
    )
    (reconstruction_loss): L2p_Loss()
    (rsr_loss): RSRLoss()
)

python

trainer = pl.Trainer(max_epochs=model.hparams.epochs, gpus=1)
trainer.fit(model, dl)

Analysing RSR Autoencoder outputs

First, let’s check the general autoencoder reconstuction.

python

model.freeze()
from torchvision.transforms import functional as tvf
from torchvision.utils import make_grid
def reconstruct(x, model):
    enc, x_hat, latent, A = model(x.view(1, -1))
    x_hat = torch.sigmoid(x_hat)
    return tvf.to_pil_image(
      make_grid([x_hat.squeeze(0).view(1, 28, 28), x])
    )

python

tvf.to_pil_image(
    make_grid([tvf.to_tensor(reconstruct(ds[i][0], model)) for i in torch.randint(0, len(ds), (8, ))], nrow=1)
)

Most of the images seem OK to me, given how simple the model is.

Now let’s look at the RSR embedding space projected into 3D using PCA.

python

import pandas as pd
rsr_embeddings = []
classes = []
errors = []
for batch in iter(DataLoader(ds, batch_size=64, shuffle=False)):
    X, cl = batch
    x = X.view(X.size(0), -1)
    enc, x_hat, latent, A = model(x)
    rsr_embeddings.append(latent)
    classes.extend(cl.numpy())
    for i in range(X.size(0)):
        rec_error = L2p_Loss()(torch.sigmoid(x_hat[i]).unsqueeze(0), x[i].unsqueeze(0))
        errors.append(float(rec_error.numpy()))
        
all_embs = torch.vstack(rsr_embeddings)
df = pd.DataFrame(all_embs.numpy(), columns=["x", "y", "z"] + [f"dim_{i}" for i in range(hparams["d"] - 3)])
df.loc[:, "class"] = classes
df.loc[:, "errors"] = errors

python

from sklearn.decomposition import PCA
pca = PCA(n_components=3)
rsr_3d = pd.DataFrame(pca.fit_transform(all_embs), columns=["x", "y", "z"])
rsr_3d.loc[:, "class"] = classes

python

import plotly.express as px
import plotly
plotly.offline.init_notebook_mode()
df = df
fig = px.scatter_3d(df, x='x', y='y', z='z',
               symbol="class", color="class", opacity=0.95)
fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))
fig.show()

Indeed, most of the 4 (the main class - inliers) reside in the center of the space, with other classes positioned further away from the center.

Unsupervised anomaly detection using RSR Autoencoder

Now let’s see if it works as the authors of the paper claim - whether the data points with high reconstruction error are really outliers.

python

df.errors.describe()

python

    count    6242.000000
    mean        1.551347
    std         0.293904
    min         0.615430
    25%         1.355003
    50%         1.512501
    75%         1.707990
    max         3.462465
    Name: errors, dtype: float64

python

df.groupby("class").errors.hist(legend=True, bins=60, figsize=(12, 12))

The reconstruction loss is centered around 1.5 with a tail spreading up to 3.46. What are the data points with the largest errors?

python

lowest_mistakes = df.sort_values(by="errors", ascending=True).head(60).loc[:, ["errors", "class"]]
highest_mistakes = df.sort_values(by="errors", ascending=False).head(60).loc[:, ["errors", "class"]]
highest_mistakes.head(10)

	errors	class
5917	3.462465	4
2201	3.120193	4
5834	3.116279	4
6017	3.109386	4
601	3.098328	4
983	3.097286	4
1312	3.091813	4
3304	3.030314	4
6015	2.998559	4
1619	2.974236	4

It seems really odd at first, as I've expected all the non-4 classes to have the largest errors. Let's carefully look at the actual results though:

python

print("Images with the highest reconsturction loss")
tvf.to_pil_image(
    make_grid([tvf.to_tensor(reconstruct(ds[i][0], model)) for i in highest_mistakes.index], nrow=6,
             pad_value=0.5)
)

Images with the highest reconsturction loss

python

print("Images with the lowest reconsturction loss")
tvf.to_pil_image(
    make_grid([tvf.to_tensor(reconstruct(ds[i][0], model)) for i in lowest_mistakes.index], nrow=6,
             pad_value=0.5)
)

Images with the lowest reconsturction loss

Now it’s much clearer - the model found the manifold that separates the images into the “normal” ones - that in this setting seem like more smooth ones with clean curves - and the outliers that are rather sharp, contain some random artifacts (which are anomalies IMHO) and usually seem slightly malformed. Given the model simplicity and the setup, it looks to me like a good starting point for the unsupervised anomaly detection.

Summary

I hope you find this post useful and it will help you to quickly get started with RSR Layer / RSRAE and unsupervised anomaly detection in PyTorch. Do not hesistate to comment and share!

Additional links & resources

Jupyter notebook for this post in Google Colab
GitHub repo with this notebook: https://github.com/marrrcin/rsrlayer-pytorch

13 features of PyTorch that you should know

Marcin Zabłocki — Sun, 21 Feb 2021 16:00:48 +0200

PyTorch has gained a lot of traction in both academia as well as in applied research in the industry. It’s a deep learning framework with great elasticity and huge number of utilities and functions to speed up the work. PyTorch’s learning curve is not that steep but implementing both efficient and clean code in it can be tricky. After using it for over 2 years, here are my top PyTorch features I wish I knew when I stared learning it.

TL;DR

13 features of PyTorch that you should know - a short list.

1. DatasetFolder

One of the first things people do when learning PyTorch is implementing their own Dataset of some kind. It’s a rookie mistake - there is no point of wasting time on writing such. Usually, datasets are either lists of data (or numpy arrays) or files on disk. It’s always better to have the data organized on a disk rather than writing custom Dataset to load the weird format that someone (or you) decided to store it with.

One of the most common data formats for classifiers is to have a directory with subfolders, representing classes, and files in those subfolders, representing examples, like the one below.

bash

folder/class_0/file1.txt
folder/class_0/file2.txt
folder/class_0/...

folder/class_1/file3.txt
folder/class_1/file4.txt

folder/class_2/file5.txt
folder/class_2/...

There is a built-in way to load this kind of datasets - does not matter if the data you have are images, text files or something else - just use DatasetFolder (https://pytorch.org/vision/0.8/datasets.html#datasetfolder). Surprisingly, this class is a part of torchvision package, not the core PyTorch. The class super comprehensive - you can filter files from the folder, load them using custom code and transform the raw files on the fly. Example:

python

from torchvision.datasets import DatasetFolder
from pathlib import Path
# I have text files in this folder
ds = DatasetFolder("/Users/marcin/Dev/tmp/my_text_dataset", 
    loader=lambda path: Path(path).read_text(),
    extensions=(".txt",), #only load .txt files
    transform=lambda text: text[:100], # only take first 100 characters
)

# Everything you need is already there
len(ds), ds.classes, ds.class_to_idx

python

(20, ['novels', 'thrillers'], {'novels': 0, 'thrillers': 1})

If you’re working with images, there is also a torchvision.datasets.ImageFolder class, which bases on the DatasetLoader and it’s preconfigured to load images.

2. Stop using `.to(device)` that much: `zeros_like` / `ones_like` etc.

I’ve read a lot fo PyTorch code from various GitHub repositories. What irritates me the most is that almost in every repo there is a lot of *.to(device) lines, which put the data from CPU to GPU or the other way around. Such lines usually occur in a lot of quick-and-dirty repos or beginner tutorials. I highly encourage of implementing such operations as little as possible and rely on the built-in PyTorch functionality that does that automatically. Putting those .to(device) lines here and there usually leads to either degraded performance or PyTorch-particioners’ favourite exception:

Expected object of device type cuda but got device type cpu

Obviously, there are some cases when you just can’t get around it, but most (if not all) of the trivial cases are covered. One of the examples is initialization of zeros/ones/random tensors that you might need in your custom loss - as the training of deep neural networks usually happens on GPU, model’s output is already on “cuda” device, but you now need to have another tensor of zeros/ones also be on “cuda” device to do operations on it. This is where *_like operations from PyTorch come in handy:

python

my_output # on any device, if it's cuda then my_zeros will also be on cuda
my_zeros = torch.zeros_like(my_output_from_model)

Under the hood, what PyTorch does is it invokes the following operation:

python

my_zeros = torch.zeros(my_output.size(), dtype=my_output.dtype, layout=my_output.layout, device=my_output.device)

so everything is set properly, and you decrease the probability of bugs in your code. Built-in operations include:

python

torch.zeros_like()
torch.ones_like()
torch.rand_like()
torch.randn_like()
torch.randint_like()
torch.empty_like()
torch.full_like()

3. Register Buffer (a.k.a `nn.Module.register_buffer`)

This is a next stop on my crusade to discourage people from using .to(device) everywhere. Sometimes your model or loss function needs to have parameters that are set upfront and are used when forward pass is invoked - for instance it can be a “weight” parameter, which scales the loss or some fixed tensor that is not changing but it’s used every time. For this kind of scenarios use nn.Module.register_buffer method, which tells PyTorch to store the values you pass to it within the module and to move those values with the module. If you initialize your module and then move it to GPU, those values will also be automatically moved. Moreover - if you save your module’s state - the buffers will also be saved!

Once registered, the values can be accessed in the forward function just like any other module’s property.

python

from torch import nn
import torch

class ModuleWithCustomValues(nn.Module):
    def __init__(self, weights, alpha):
        super().__init__()
        self.register_buffer("weights", torch.tensor(weights))
        self.register_buffer("alpha", torch.tensor(alpha))
    
    def forward(self, x):
        return x * self.weights + self.alpha

m = ModuleWithCustomValues(
    weights=[1.0, 2.0], alpha=1e-4
)
m(torch.tensor([1.23, 4.56]))

python

tensor([1.2301, 9.1201])

4. Built-in `Identity()`

Sometimes when you play with transfer learning you will need to replace some layers with 1:1 mapping, which boils down to implemeting a nn.Module with the only purpose or returning the input value. PyTorch has this class built-in: https://pytorch.org/docs/1.7.1/generated/torch.nn.Identity.html?highlight=identity#identity.

Example - you want to get image representation from a pre-trained ResNet50 just before the classification layer. Here is how to do this:

python

from torchvision.models import resnet50
model = resnet50(pretrained=True)
model.fc = nn.Identity()
last_layer_output = model(torch.rand((1, 3, 224, 224)))
last_layer_output.shape

python

torch.Size([1, 2048])

5. Pairwise distances: `torch.cdist`

The next time you will encounter a problem of calculating all-pairs euclidean (or in general: a p-norm) distance between two tensors, remember about torch.cdist. It does exactly that and also automatically uses matrix multiplication when euclidean distance is used, giving a performance boost.
Docs: https://pytorch.org/docs/1.7.1/generated/torch.cdist.html#torch-cdist

python

points1 = torch.tensor([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]])
points2 = torch.tensor([[0.0, 0.0], [-1.0, -1.0], [-2.0, -2.0], [-3.0, -3.0]]) # batches don't have to be equal
torch.cdist(points1, points2, p=2.0)

python

tensor([[0.0000, 1.4142, 2.8284, 4.2426],
        [1.4142, 2.8284, 4.2426, 5.6569],
        [2.8284, 4.2426, 5.6569, 7.0711]])

Performance of without and with matrix multiplications - it’s over 2x faster on my machine when using mm.

python

%%timeit
points1 = torch.rand((512, 2))
points2 = torch.rand((512, 2))
torch.cdist(points1, points2, p=2.0, compute_mode="donot_use_mm_for_euclid_dist")

867 µs ± 142 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

python

%%timeit
points1 = torch.rand((512, 2))
points2 = torch.rand((512, 2))
torch.cdist(points1, points2, p=2.0)

417 µs ± 52.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

6. Cosine similarity: `F.cosine_similarity`

Staying within the same topic as in the last point - calculating distances - euclidean distance is not always the thing you need. When working with vectors, usually the cosine similarity is the metric of choice. PyTorch has a built-in implementation of cosine similarity too.

python

import torch.nn.functional as F
vector1 = torch.tensor([0.0, 1.0])
vector2 = torch.tensor([0.05, 1.0])
print(F.cosine_similarity(vector1, vector2, dim=0))
vector3 = torch.tensor([0.0, -1.0])
print(F.cosine_similarity(vector1, vector3, dim=0))

python

tensor(0.9988)
tensor(-1.)

Batched cosine similarity in PyTorch

python

import torch.nn.functional as F
batch_of_vectors = torch.rand((4, 64))
similarity_matrix = F.cosine_similarity(batch_of_vectors.unsqueeze(1), batch_of_vectors.unsqueeze(0), dim=2)
similarity_matrix

python

tensor([[1.0000, 0.6922, 0.6480, 0.6789],
        [0.6922, 1.0000, 0.7143, 0.7172],
        [0.6480, 0.7143, 1.0000, 0.7312],
        [0.6789, 0.7172, 0.7312, 1.0000]])

7. Normalizing vectors: `F.normalize`

Last point still loosely connected to vectors and distances is normalization: usually it’s used to improve the stability of calculation, by changing the magnitude of the vector. The most commonly used normalization is L2 one and can be applied in PyTorch as follows:

python

vector = torch.tensor([99.0, -512.0, 123.0, 0.1, 6.66])
normalized_vector = F.normalize(vector, p=2.0, dim=0)
normalized_vector

python

tensor([ 1.8476e-01, -9.5552e-01,  2.2955e-01,  1.8662e-04,  1.2429e-02])

The old way of doing normalization in PyTorch was:

python

vector = torch.tensor([99.0, -512.0, 123.0, 0.1, 6.66])
normalized_vector = vector / torch.norm(vector, p=2.0)
normalized_vector

python

tensor([ 1.8476e-01, -9.5552e-01,  2.2955e-01,  1.8662e-04,  1.2429e-02])

Batched L2 normalization in PyTorch

python

batch_of_vectors = torch.rand((4, 64))
normalized_batch_of_vectors = F.normalize(batch_of_vectors, p=2.0, dim=1)
normalized_batch_of_vectors.shape, torch.norm(normalized_batch_of_vectors, dim=1) # all vectors will have length of 1.0

python

(torch.Size([4, 64]), tensor([1.0000, 1.0000, 1.0000, 1.0000]))

8. Linear layer + chunking trick (`torch.chunk`)

This is a creative trick I’ve found recently here. Let’s say you want to map your input into N different linear projections. You can do this by creating N nn.Linear layers and doing forward pass with all N of them OR you can create a single linear layer, do one forward pass and just chunk the output into N pieces. This method usually leads to higher performance, so it’s a nice trick to remember.

python

d = 1024
batch = torch.rand((8, d))
layers = nn.Linear(d, 128, bias=False), nn.Linear(d, 128, bias=False), nn.Linear(d, 128, bias=False)
one_layer = nn.Linear(d, 128 * 3, bias=False)

python

%%timeit
o1 = layers[0](batch)
o2 = layers[1](batch)
o3 = layers[2](batch)

289 µs ± 30.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

python

%%timeit
o1, o2, o3 = torch.chunk(one_layer(batch), 3, dim=1)

202 µs ± 8.09 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

9. Masked select (`torch.masked_select`)

Sometimes you will need to do your calculations only on some portion of the input tensors. To give you an example: you want to calculate loss only on the tensors that fulfuill some condition. In order to do so, just use torch.masked_select - note that this operation can be used when gradient is required too.

python

data = torch.rand((3, 3)).requires_grad_()
print(data)
mask = data > data.mean()
print(mask)
torch.masked_select(data, mask)

python

tensor([[0.0582, 0.7170, 0.7713],
        [0.9458, 0.2597, 0.6711],
        [0.2828, 0.2232, 0.1981]], requires_grad=True)
tensor([[False,  True,  True],
        [ True, False,  True],
        [False, False, False]])
tensor([0.7170, 0.7713, 0.9458, 0.6711], grad_fn=<MaskedSelectBackward>)

Applying mask on tensor directly

Similar behaviour can be achiveded by using mask as an “indexer” of the input tensor.

python

data[mask]

python

tensor([0.7170, 0.7713, 0.9458, 0.6711], grad_fn=<IndexBackward>)

Sometimes, a desirable solution is to fill all False values from mask with zeros, which can be done like this:

python

data * mask

python

tensor([[0.0000, 0.7170, 0.7713],
        [0.9458, 0.0000, 0.6711],
        [0.0000, 0.0000, 0.0000]], grad_fn=<MulBackward0>)

10. Conditional tensors with `torch.where`

This functions comes in handy when you want to combine two tensors with the condition - if it’s true, then take element from first tensor, if it’s false - take from the second one.
Docs: https://pytorch.org/docs/1.7.1/generated/torch.where.html#torch.where

python

x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], requires_grad=True)
y = -x
condition_or_mask = x <= 3.0
torch.where(condition_or_mask, x, y)

python

tensor([ 1.,  2.,  3., -4., -5.], grad_fn=<SWhereBackward>)

11. Filling tensors with values at given positions (`Tensor.scatter`)

Use case for this function is the following - you want to fill one tensor with values from another tensor at given indices. It’s simpler to understand on a 1D tensor, so I will show it first and then proceed to the more advanced example.

python

data = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1])
values = torch.tensor([-1, -2, -3, -4, -5])
data.scatter(0, index, values)

python

tensor([-1, -2,  3,  4,  5])

The above example is simple, but now see what happens if you change index to index = torch.tensor([0, 1, 4])

python

data = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1, 4])
values = torch.tensor([-1, -2, -3, -4, -5])
data.scatter(0, index, values)

python

tensor([-1, -2,  3,  4, -3])

It’s counter-intuitive why the last value is now -3, right? This is the clou of the PyTorch scatter function. The index variable says at which position from data tensor should the i-th value from the values tensor be placed. I hope the below plain python equivalent of this operation will shed more light:

python

data_orig = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1, 4])
values = torch.tensor([-1, -2, -3, -4, -5])
scattered = data_orig.scatter(0, index, values)

data = data_orig.clone()
for idx_in_values, where_to_put_the_value in enumerate(index):
    what_value_to_put = values[idx_in_values]
    data[where_to_put_the_value] = what_value_to_put
data, scattered

python

(tensor([-1, -2,  3,  4, -3]), tensor([-1, -2,  3,  4, -3]))

PyTorch scatter example on 2D data

Always remember that shape of the index is related to the shape of the values and values within index correspond to the positions in data.

python

data = torch.zeros((4, 4)).float()
index = torch.tensor([
    [0, 1],
    [2, 3],
    [0, 3],
    [1, 2]
])
values = torch.arange(1, 9).float().view(4, 2)
values, data.scatter(1, index, values)

python

(tensor([[1., 2.],
        [3., 4.],
        [5., 6.],
        [7., 8.]]),
tensor([[1., 2., 0., 0.],
        [0., 0., 3., 4.],
        [5., 0., 0., 6.],
        [0., 7., 8., 0.]]))

12. Image interpolation within network (`F.interpolate`)

When I was learning PyTorch it surprised me that you can actually resize the image (or any intermediate tensor) within the forward pass and maintain the gradient flow. This method is especially useful when working with CNNs and GANs.

python

# image from https://commons.wikimedia.org/wiki/File:A_female_British_Shorthair_at_the_age_of_20_months.jpg
img = Image.open("./cat.jpg")
img

python

to_pil_image(
    F.interpolate(to_tensor(img).unsqueeze(0),  # batch of size 1
                  mode="bilinear", 
                  scale_factor=2.0, 
                  align_corners=False).squeeze(0) # remove batch dimension
)

See how the gradient flow is preserved:

python

F.interpolate(to_tensor(img).unsqueeze(0).requires_grad_(),
                  mode="bicubic", 
                  scale_factor=2.0, 
                  align_corners=False)

python

tensor([[[[0.9216, 0.9216, 0.9216,  ..., 0.8361, 0.8272, 0.8219],
    [0.9214, 0.9214, 0.9214,  ..., 0.8361, 0.8272, 0.8219],
    [0.9212, 0.9212, 0.9212,  ..., 0.8361, 0.8272, 0.8219],
    ...,
    [0.9098, 0.9098, 0.9098,  ..., 0.3592, 0.3486, 0.3421],
    [0.9098, 0.9098, 0.9098,  ..., 0.3566, 0.3463, 0.3400],
    [0.9098, 0.9098, 0.9098,  ..., 0.3550, 0.3449, 0.3387]],

    [[0.6627, 0.6627, 0.6627,  ..., 0.5380, 0.5292, 0.5238],
    [0.6626, 0.6626, 0.6626,  ..., 0.5380, 0.5292, 0.5238],
    [0.6623, 0.6623, 0.6623,  ..., 0.5380, 0.5292, 0.5238],
    ...,
    [0.6196, 0.6196, 0.6196,  ..., 0.3631, 0.3525, 0.3461],
    [0.6196, 0.6196, 0.6196,  ..., 0.3605, 0.3502, 0.3439],
    [0.6196, 0.6196, 0.6196,  ..., 0.3589, 0.3488, 0.3426]],

    [[0.4353, 0.4353, 0.4353,  ..., 0.1913, 0.1835, 0.1787],
    [0.4352, 0.4352, 0.4352,  ..., 0.1913, 0.1835, 0.1787],
    [0.4349, 0.4349, 0.4349,  ..., 0.1913, 0.1835, 0.1787],
    ...,
    [0.3333, 0.3333, 0.3333,  ..., 0.3827, 0.3721, 0.3657],
    [0.3333, 0.3333, 0.3333,  ..., 0.3801, 0.3698, 0.3635],
    [0.3333, 0.3333, 0.3333,  ..., 0.3785, 0.3684, 0.3622]]]],
grad_fn=<UpsampleBicubic2DBackward1>)

13. Making grid of images (`torchvision.utils.make_grid`)

There is no need to copy-paste the code using matplotlib or some external libraries in order to display a grid of images when working with PyTorch and torchvision. Just use torchvision.utils.make_grid!
Docs: https://pytorch.org/vision/0.8/utils.html

python

from torchvision.utils import make_grid
from torchvision.transforms.functional import to_tensor, to_pil_image
from PIL import Image
img = Image.open("./cat.jpg")
to_pil_image(
    make_grid(
        [to_tensor(i) for i in [img, img, img]],
         nrow=2, # number of images in single row
         padding=5 # "frame" size
     )
)

Summary

I hope my post will help you in your journey to master PyTorch. Please share if you like it and don’t hesitate to ask questions in comments!

Using ResNet for MNIST in PyTorch 1.7

Marcin Zabłocki — Sat, 30 Jan 2021 16:00:00 +0000

This short post is a refreshed version of my early-2019 post about adjusting ResNet architecture for use with well known MNIST dataset. The goal of this post is to provide refreshed overview on this process for the beginners. Treat is a tutorial how to train a MNIST digits classifier using PyTorch 1.7 and Torchvision.

TL;DR

Tutorial on how to train ResNet for MNIST using PyTorch, updated for 2021. Link to google colab at the bottom.

Use ResNet from torchvision

If you’re new to PyTorch, then you should know that there is an official PyTorch library dedicated for computer vision problems - Torchvision. It provides stable and well tested implementation of various network architectures, including: ResNets in many variants (ResNet-18, ResNet-34, ResNet-51 and so on), ShuffleNets, MobileNets as well as many other ResNet-related ones, that improve the original architecture e.g. ResNeXt.

Torchvision also weights for those networks that are trained on ImageNet dataset and are ready to use for transfer learning. In this tutorial, we will be training ResNet-18 from scratch, as digit images from MNIST dataset are really different from the images included in the ImageNet.

python

from torchvision.models import resnet18
from torch import nn
model = resnet18(num_classes=10) # MNIST has 10 classes

The overall model architecture looks like this:

python

ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  # Skipped for brevity
  (layer1): ...
  (layer2): ...
  (layer3): ...
  (layer4): ...
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=512, out_features=10, bias=True)
)

In order to adapt this architecture for MNIST there is one more thing required - input layer needs to accept single channel instead of 3 (MNIST images are single-channel = grayscale, whereas ImageNet are 3-channels = RGB).

python

model.conv1 = Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

And it’s done! No more modifications are required.

Load the dataset

Besides the models, Torchvision also gives you access to a lot of well known computer vision datasets, including MNIST - it will automatically download the dataset for you and provide you with PyTorch’s Dataset abstraction that it’s ready to use.

python

train_ds = MNIST("mnist", train=True, download=True)
test_ds = MNIST("mnist", train=False, download=True)
train_dl = DataLoader(train_ds, batch_size=64, shuffle=True)
test_dl = DataLoader(test_ds, batch_size=64)

Do not write training loops in PyTorch

In the previous post, I’ve showed the training loop boiler plate for PyTorch. Back in early 2019 there were no well established libraries for trianing PyTorch models. Fortunately, in the early 2021 there are plenty. One of them is PyTorch Lightning. The minimal code is as follows:

python

import pytorch_lightning as pl
from pytorch_lightning.core.decorators import auto_move_data

class ResNetMNIST(pl.LightningModule):
  def __init__(self):
    super().__init__()
    # define model and loss
    self.model = resnet18(num_classes=10)
    self.model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    self.loss = nn.CrossEntropyLoss()

  @auto_move_data # this decorator automatically handles moving your tensors to GPU if required
  def forward(self, x):
    return self.model(x)
  
  def training_step(self, batch, batch_no):
    # implement single training step
    x, y = batch
    logits = self(x)
    loss = self.loss(logits, y)
    return loss
  
  def configure_optimizers(self):
    # choose your optimizer
    return torch.optim.RMSprop(self.parameters(), lr=0.005)

python

model = ResNetMNIST()
trainer = pl.Trainer(
    gpus=1, # use one GPU
    max_epochs=1, # set number of epochs
    progress_bar_refresh_rate=20 # set to >= 20 if running in Google Colab
)
trainer.fit(model, train_dl)

bash

  | Name  | Type             | Params
-------------------------------------------
0 | model | ResNet           | 11.2 M
1 | loss  | CrossEntropyLoss | 0     
-------------------------------------------
Epoch 0: 100% 938/938 [00:21<00:00, 43.77it/s, loss=0.116]

And it’s done! Now, save the model so that it can be used later for inference.

python

trainer.save_checkpoint("resnet18_mnist.pt")

Inference with PyTorch Lightning

The infernce code is pretty straightforward. I’ve only added tqdm so that it reports the progress.

python

from tqdm.autonotebook import tqdm

def get_prediction(x, model: pl.LightningModule):
  model.freeze() # prepares model for predicting
  probabilities = torch.softmax(model(x), dim=1)
  predicted_class = torch.argmax(probabilities, dim=1)
  return predicted_class, probabilities

inference_model = ResNetMNIST.load_from_checkpoint("resnet18_mnist.pt", map_location="cuda")

Loop through any data…

python

true_y, pred_y = [], []
for batch in tqdm(iter(test_dl), total=len(test_dl)):
  x, y = batch
  true_y.extend(y)
  preds, probs = get_prediction(x, inference_model)
  pred_y.extend(preds.cpu())

…and get the final classification report:

python

from sklearn.metrics import classification_report
print(classification_report(true_y, pred_y, digits=3))

              precision    recall  f1-score   support

           0      0.963     0.996     0.979       980
           1      0.991     0.993     0.992      1135
           2      0.974     0.980     0.977      1032
           3      0.994     0.975     0.985      1010
           4      0.990     0.981     0.985       982
           5      0.974     0.976     0.975       892
           6      0.975     0.981     0.978       958
           7      0.966     0.984     0.975      1028
           8      0.977     0.963     0.970       974
           9      0.988     0.960     0.974      1009

    accuracy                          0.979     10000
   macro avg      0.979     0.979     0.979     10000
weighted avg      0.979     0.979     0.979     10000

Summary

I hope you find this post useful for your first steps with neural networks and PyTorch. Do not hesitate to comment and / or ask any questions.

Additional links & resources

Jupyter notebook for this post in Google Colab
GitHub repo with this notebook: https://github.com/marrrcin/pytorch-resnet-mnist

Understanding & implementing SimCLR in PyTorch - an ELI5 guide

Marcin Zabłocki — Mon, 27 Jul 2020 14:57:48 +0200

Transfer learning and pre-training schemas for both NLP and Computer Vision have gained a lot of attention in the last months. Research showed that carefully designed unsupervised/self-supervised training can produce high quality base models and embeddings that greatly decrease the amount of data needed to obtain good classification models downstream. This approach becomes more and more important as the companies collect a lot of data from which only a fraction can be labelled by humans - either due to large cost of labelling process or due to some time constraints.

Here I explore SimCLR pre-training framework proposed by Google in this arxiv paper. I will explain the SimCLR and its contrastive loss function step by step, starting from naive implementation in PyTorch, followed by faster, vectorized one. Then I will show how to use SimCLR’s pre-training routine to first build image embeddings using EfficientNet network architecture and finally I will show how to build a classifier on top of it.

TL;DR

This post covers:

understanding the SimCLR framework with code samples in PyTorch
from scratch explanation & implementation of SimCLR’s loss function (NT-Xent) in PyTorch
pre-training image embeddings using EfficientNet architecture
training classifier by using transfer learning from the pre-trained embeddings

Understanding SimCLR framework

In general, SimCLR is a simple framework for contrastive learning of visual representations. It’s not any new framework for deep learning, it’s a set of fixed steps that one should follow in order to train image embeddings of good quality. I drew a schema which explains the flow and the whole representation learning process (click to zoom).

The flow is as follows (from left to right):

Take an input image
Prepare 2 random augmentations on the image, including: rotations, hue/saturation/brightness changes, zooming, cropping, etc. The range of augmentations as well as analysis which ones work best are discussed in detail in the paper.
Run a deep neural network (preferably convolutional one, like ResNet50) to obtain image representations (embeddings) for those augmented images.
Run a small fully connected linear neural network to project embeddings into another vector space.
Calculate contrastive loss and run backpropagation through both networks. Contrastive loss decreases when projections coming from the same image are similar. The similarity between projections can be arbitrary, here I will use cosine similarity, same as in the paper.

Contrastive loss function

Theory behind contrastive loss function

One can reason about contrastive loss function form two angles:

Contrastive loss decreases when projections of augmented images coming from the same input image are similar.
For two augmented images: (i), (j) (coming from the same input image - I will call them “positive” pair later on), the contrastive loss for (i) tries to identify (j) among other images (“negative” examples) that are in the same batch .

Formal definition of the loss for pair of positive examples (i) and (j) defined as:

The final loss is an arithmetic mean of the losses for all positive pairs in the batch:

(keep in mind that the indexing in l(2k-1, 2k) + l(2k, 2k-1) is purely dependent on how you implement the loss - I find it easier to understand when I reason about them as l(i,j) + l(j,i)).

Contrastive loss function - implementation in PyTorch, ELI5 version

It’s much easier to implement the loss function without vectorization first and then follow up with the vectorization phase.

python

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

python

class ContrastiveLossELI5(nn.Module):
    def __init__(self, batch_size, temperature=0.5, verbose=True):
        super().__init__()
        self.batch_size = batch_size
        self.register_buffer("temperature", torch.tensor(temperature))
        self.verbose = verbose
            
    def forward(self, emb_i, emb_j):
        """
        emb_i and emb_j are batches of embeddings, where corresponding indices are pairs
        z_i, z_j as per SimCLR paper
        """
        z_i = F.normalize(emb_i, dim=1)
        z_j = F.normalize(emb_j, dim=1)

        representations = torch.cat([z_i, z_j], dim=0)
        similarity_matrix = F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)
        if self.verbose: print("Similarity matrix\n", similarity_matrix, "\n")
            
        def l_ij(i, j):
            z_i_, z_j_ = representations[i], representations[j]
            sim_i_j = similarity_matrix[i, j]
            if self.verbose: print(f"sim({i}, {j})={sim_i_j}")
                
            numerator = torch.exp(sim_i_j / self.temperature)
            one_for_not_i = torch.ones((2 * self.batch_size, )).to(emb_i.device).scatter_(0, torch.tensor([i]), 0.0)
            if self.verbose: print(f"1{{k!={i}}}",one_for_not_i)
            
            denominator = torch.sum(
                one_for_not_i * torch.exp(similarity_matrix[i, :] / self.temperature)
            )    
            if self.verbose: print("Denominator", denominator)
                
            loss_ij = -torch.log(numerator / denominator)
            if self.verbose: print(f"loss({i},{j})={loss_ij}\n")
                
            return loss_ij.squeeze(0)

        N = self.batch_size
        loss = 0.0
        for k in range(0, N):
            loss += l_ij(k, k + N) + l_ij(k + N, k)
        return 1.0 / (2*N) * loss

Explanation

Contrastive loss needs to know the batch size and temperature (scaling) parameter. You can find details about setting the optimal temperature parameter in the paper.

My implementation of the forward of the contrastive loss takes two parameters. First one will be a batch projection of images after first augmentation, the second will be a batch projection of images after second augmentation.

Projections needs to be normalized first, hence:

python

z_i = F.normalize(emb_i, dim=1)
z_j = F.normalize(emb_j, dim=1)

All representations are concatenated together in order to efficiently calculate cosine similarities between each image pair.

python

representations = torch.cat([z_i, z_j], dim=0)
similarity_matrix = F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)

Next is the naive implementation of l(i,j) for clarity and easiness of understanding. The code bellow almost directly implements the equation:

python

def l_ij(i, j):
    z_i_, z_j_ = representations[i], representations[j]
    sim_i_j = similarity_matrix[i, j]
    numerator = torch.exp(sim_i_j / self.temperature)
    one_for_not_i = torch.ones((2 * self.batch_size, )).to(emb_i.device).scatter_(0, torch.tensor([i]), 0.0)
    denominator = torch.sum(
        one_for_not_i * torch.exp(similarity_matrix[i, :] / self.temperature)
    )

    loss_ij = -torch.log(numerator / denominator)
    return loss_ij.squeeze(0)

Then, the final loss for the batch is computed as an arithmetic mean of all combinations of positive examples:

python

N = self.batch_size
loss = 0.0
for k in range(0, N):
    loss += l_ij(k, k + N) + l_ij(k + N, k)
return 1.0 / (2*N) * loss

Now, let’s run it with verbose mode to see what’s inside.

python

I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]])
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]])
loss_eli5 = ContrastiveLossELI5(batch_size=3, temperature=1.0, verbose=True)
loss_eli5(I, J)

bash

Similarity matrix
    tensor([[ 1.0000, -0.1240,  0.9648,  0.8944, -0.0948,  0.9679],
        [-0.1240,  1.0000, -0.3807,  0.3328,  0.9996, -0.3694],
        [ 0.9648, -0.3807,  1.0000,  0.7452, -0.3534,  0.9999],
        [ 0.8944,  0.3328,  0.7452,  1.0000,  0.3604,  0.7533],
        [-0.0948,  0.9996, -0.3534,  0.3604,  1.0000, -0.3419],
        [ 0.9679, -0.3694,  0.9999,  0.7533, -0.3419,  1.0000]]) 

sim(0, 3)=0.8944272398948669
1{k!=0} tensor([0., 1., 1., 1., 1., 1.])
Denominator tensor(9.4954)
loss(0,3)=1.3563847541809082

sim(3, 0)=0.8944272398948669
1{k!=3} tensor([1., 1., 1., 0., 1., 1.])
Denominator tensor(9.5058)
loss(3,0)=1.357473373413086

sim(1, 4)=0.9995677471160889
1{k!=1} tensor([1., 0., 1., 1., 1., 1.])
Denominator tensor(6.3699)
loss(1,4)=0.8520082831382751

sim(4, 1)=0.9995677471160889
1{k!=4} tensor([1., 1., 1., 1., 0., 1.])
Denominator tensor(6.4733)
loss(4,1)=0.8681114912033081

sim(2, 5)=0.9999250769615173
1{k!=2} tensor([1., 1., 0., 1., 1., 1.])
Denominator tensor(8.8348)
loss(2,5)=1.1787779331207275

sim(5, 2)=0.9999250769615173
1{k!=5} tensor([1., 1., 1., 1., 1., 0.])
Denominator tensor(8.8762)
loss(5,2)=1.1834462881088257

tensor(1.1327)

A few things happened there, but by going back and forward between the verbose logs and the equation, everything should become clear. The indexing jumps by batch size (first l(0,3), l(3,0) then l(1,4), l(4,1) because of the way the similarity matrix was constructed. First row of the similarity_matrix is:

python

[ 1.0000, -0.1240,  0.9648,  0.8944, -0.0948,  0.9679]

Remember the input:

python

I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]])
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]])

Now:

1.0000 is the cosine similarity between I[0] and I[0] ([1.0, 2.0] and [1.0, 2.0])

-0.1240 is the cosine similarity between I[0] and I[1] ([1.0, 2.0] and [3.0, -2.0])

-0.0948 is the cosine similarity between I[0] and J[2] ([1.0, 2.0] and [2.8, -1.75])

… and so on.

Let’s see if the loss decreases if the similarity between first image projection increases:

python

I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]])
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]])
J = torch.tensor([[1.0, 1.75], [2.8, -1.75], [1.0, 4.7]]) # note the change
ContrastiveLossELI5(3, 1.0, verbose=False)(I, J)

bash

tensor(1.0996)

Indeed the loss decreases! Now I will follow up with the vectorized implementation.

Contrastive loss function - implementation in PyTorch, vectorized version

Performance of naive implementation is really poor (mostly due to manual loop), see the results:

python

contrastive_loss_eli5 = ContrastiveLossELI5(3, 1.0, verbose=False)
I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]], requires_grad=True)
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]], requires_grad=True)

python

%%timeit
contrastive_loss_eli5(I, J)
838 µs ± 23.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Once I understood the internals of the loss, it’s easy to vectorize it and remove the manual loop:

python

class ContrastiveLoss(nn.Module):
    def __init__(self, batch_size, temperature=0.5):
        super().__init__()
        self.batch_size = batch_size
        self.register_buffer("temperature", torch.tensor(temperature))
        self.register_buffer("negatives_mask", (~torch.eye(batch_size * 2, batch_size * 2, dtype=bool)).float())
            
    def forward(self, emb_i, emb_j):
        """
        emb_i and emb_j are batches of embeddings, where corresponding indices are pairs
        z_i, z_j as per SimCLR paper
        """
        z_i = F.normalize(emb_i, dim=1)
        z_j = F.normalize(emb_j, dim=1)

        representations = torch.cat([z_i, z_j], dim=0)
        similarity_matrix = F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)
        
        sim_ij = torch.diag(similarity_matrix, self.batch_size)
        sim_ji = torch.diag(similarity_matrix, -self.batch_size)
        positives = torch.cat([sim_ij, sim_ji], dim=0)
        
        nominator = torch.exp(positives / self.temperature)
        denominator = self.negatives_mask * torch.exp(similarity_matrix / self.temperature)
    
        loss_partial = -torch.log(nominator / torch.sum(denominator, dim=1))
        loss = torch.sum(loss_partial) / (2 * self.batch_size)
        return loss

python

contrastive_loss = ContrastiveLoss(3, 1.0)
contrastive_loss(I, J).item() - contrastive_loss_eli5(I, J).item()
0.0

The difference should be zero or close to zero (< 1e-6 due to fp arithmetics). Performance comparison:

python

I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]], requires_grad=True)
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]], requires_grad=True)

python

%%timeit
contrastive_loss_eli5(I, J)
918 µs ± 60.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

python

%%timeit
contrastive_loss(I, J)
272 µs ± 9.18 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Almost 4x improvement, it works.

Pre-training image embeddings using SimCLR with EfficientNet

Once the loss function is established and understood it’s time to make a good use of it. I will pre-train image embeddings using EfficientNet architecture, following the SimCLR framework. For convenience, I’ve implemented a few utility functions and classes that I will explain briefly below. The training code is structured using PyTorch-Lightning.

I will use a great EfficientNet [https://arxiv.org/abs/1905.11946] implementation by Luke Melas-Kyriazi from GitHub, already pre-trained on ImageNet (transfer learning inception!). The dataset I choose is STL10 (from torchvision) as it contains both training and unlabeled splits for unsupervised / self-supervised learning tasks.

My goal here is to demonstrate the whole SimCLR flow from start to the end. I had no intent to reach new SOTA with the presented configuration.

Utility functions for image augmentations

Training with SimCLR produces good image embeddings that are not affected by image transformations - this is because during training, various data augmentations are done to force the network to understand the contents of the images regardless of i.e color of the image or the position where the object in the image is placed. SimCLR’s authors say that composition of data augmentations plays a critical role in defining effective predictive tasks and also contrastive learning needs stronger data augmentation than supervised learning. To sum this up: when pre-training the image embeddings it’s good to make this task difficult for the network to learn by strongly augmenting the images in order to generalize better afterwards.

I strongly advise to read the SimCLR’s paper and it’s appendix as they did ablation studies on which data augmentations bring the best effects on the embeddings.

To keep this blog-post simple to go through, I will mostly use built-in Torchvision’s data augmentations , with one additional one - random resized rotation.

python

def random_rotate(image):
    if random.random() > 0.5:
        return tvf.rotate(image, angle=random.choice((0, 90, 180, 270)))
    return image

python

class ResizedRotation():
    def __init__(self, angle, output_size=(96, 96)):
        self.angle = angle
        self.output_size = output_size
        
    def angle_to_rad(self, ang): return np.pi * ang / 180.0
        
    def __call__(self, image):
        w, h = image.size
        new_h = int(np.abs(w * np.sin(self.angle_to_rad(90 - self.angle))) + np.abs(h * np.sin(self.angle_to_rad(self.angle))))
        new_w = int(np.abs(h * np.sin(self.angle_to_rad(90 - self.angle))) + np.abs(w * np.sin(self.angle_to_rad(self.angle))))
        img = tvf.resize(image, (new_w, new_h))
        img = tvf.rotate(img, self.angle)
        img = tvf.center_crop(img, self.output_size)
        return img

python

class WrapWithRandomParams():
    def __init__(self, constructor, ranges):
        self.constructor = constructor
        self.ranges = ranges
    
    def __call__(self, image):
        randoms = [float(np.random.uniform(low, high)) for _, (low, high) in zip(range(len(self.ranges)), self.ranges)]
        return self.constructor(*randoms)(image)

python

from torchvision.datasets import STL10
import torchvision.transforms.functional as tvf
from torchvision import transforms
import numpy as np

A brief look on the transformation results:

python

stl10_unlabeled = STL10(".", split="unlabeled", download=True)

python

idx = 123

python

random_resized_rotation = WrapWithRandomParams(lambda angle: ResizedRotation(angle), [(0.0, 360.0)])

python

random_resized_rotation(tvf.resize(stl10_unlabeled[idx][0], (96, 96)))

Automatic data augmentation wrapper

Here I’ve also implemented a utility dataset wrapper that automatically applies random data augmentations every time when the image is retrieved. It can be easily used with any image dataset as long as it follows the simple interface of returning tuple with (PIL Image, anything). This wrapper can be set to return a deterministic transformation with the debug flag set to True. Note that there is a preprocess step that applies ImageNet-originated data standardization as I’m using already pre-trained EfficientNet.

python

from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler, SequentialSampler
import random

python

class PretrainingDatasetWrapper(Dataset):
    def __init__(self, ds: Dataset, target_size=(96, 96), debug=False):
        super().__init__()
        self.ds = ds
        self.debug = debug
        self.target_size = target_size
        if debug:
            print("DATASET IN DEBUG MODE")
        
        # I will be using network pre-trained on ImageNet first, which uses this normalization.
        # Remove this, if you're training from scratch or apply different transformations accordingly
        self.preprocess = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        
        random_resized_rotation = WrapWithRandomParams(lambda angle: ResizedRotation(angle, target_size), [(0.0, 360.0)])
        self.randomize = transforms.Compose([
            transforms.RandomResizedCrop(target_size, scale=(1/3, 1.0), ratio=(0.3, 2.0)),
            transforms.RandomChoice([
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.Lambda(random_rotate)
            ]),
            transforms.RandomApply([
                random_resized_rotation
            ], p=0.33),
            transforms.RandomApply([
                transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.2)
            ], p=0.8),
            transforms.RandomGrayscale(p=0.2)
        ])
    
    def __len__(self): return len(self.ds)
    
    def __getitem_internal__(self, idx, preprocess=True):
        this_image_raw, _ = self.ds[idx]
        
        if self.debug:
            random.seed(idx)
            t1 = self.randomize(this_image_raw)
            random.seed(idx + 1)
            t2 = self.randomize(this_image_raw)
        else:
            t1 = self.randomize(this_image_raw)
            t2 = self.randomize(this_image_raw)
        
        if preprocess:
            t1 = self.preprocess(t1)
            t2 = self.preprocess(t2)
        else:
            t1 = transforms.ToTensor()(t1)
            t2 = transforms.ToTensor()(t2)

        return (t1, t2), torch.tensor(0)

    def __getitem__(self, idx):
        return self.__getitem_internal__(idx, True)
    
    def raw(self, idx):
        return self.__getitem_internal__(idx, False)

python

ds = PretrainingDatasetWrapper(stl10_unlabeled, debug=False)

python

tvf.to_pil_image(ds[idx][0][0])

python

tvf.to_pil_image(ds[idx][0][1])

python

tvf.to_pil_image(ds.raw(idx)[0][1])

python

tvf.to_pil_image(ds.raw(idx)[0][0])

SimCLR neural network for embeddings

Here I define the ImageEmbedding neural network which is based on EfficientNet-b0 architecture. I swap out the last layer of pre-trained EfficientNet with identity function and add projection for image embeddings on top of it (following the SimCLR paper) with Linear-ReLU-Linear layers. It was shown in the paper that the non-linear projection head (i.e Linear-ReLU-Linear) improves the quality of the embeddings.

python

from efficientnet_pytorch import EfficientNet

python

class ImageEmbedding(nn.Module):       
    class Identity(nn.Module):
        def __init__(self): super().__init__()

        def forward(self, x):
            return x
    
        
    def __init__(self, embedding_size=1024):
        super().__init__()
        
        base_model = EfficientNet.from_pretrained("efficientnet-b0")
        internal_embedding_size = base_model._fc.in_features
        base_model._fc = ImageEmbedding.Identity()
        
        self.embedding = base_model
        
        self.projection = nn.Sequential(
            nn.Linear(in_features=internal_embedding_size, out_features=embedding_size),
            nn.ReLU(),
            nn.Linear(in_features=embedding_size, out_features=embedding_size)
        )

    def calculate_embedding(self, image):
        return self.embedding(image)

    def forward(self, X):
        image = X
        embedding = self.calculate_embedding(image)
        projection = self.projection(embedding)
        return embedding, projection

Next is the implementation of PyTorch-Lightning-based training module that orchestrates everything together:

hyper-parameters handling
SimCLR ImageEmbedding network
STL10 dataset
optimizer
forward step

As the PretrainingDatasetWrapper I’ve implemented returns tuple of: (Image1, Image2), dummy class, the forward step for this module is straightforward - it needs to produce two batches of embeddings and calculate contrastive loss function:

python

(X, Y), y = batch
embX, projectionX = self.forward(X)
embY, projectionY = self.forward(Y)
loss = self.loss(projectionX, projectionY)

python

from torch.multiprocessing import cpu_count
from torch.optim import RMSprop
import pytorch_lightning as pl

python

class ImageEmbeddingModule(pl.LightningModule):
    def __init__(self, hparams):
        hparams = Namespace(**hparams) if isinstance(hparams, dict) else hparams
        super().__init__()
        self.hparams = hparams
        self.model = ImageEmbedding()
        self.loss = ContrastiveLoss(hparams.batch_size)
    
    def total_steps(self):
        return len(self.train_dataloader()) // self.hparams.epochs
    
    def train_dataloader(self):
        return DataLoader(PretrainingDatasetWrapper(stl10_unlabeled, 
                                             debug=getattr(self.hparams, "debug", False)),
                          batch_size=self.hparams.batch_size, 
                          num_workers=cpu_count(),
                          sampler=SubsetRandomSampler(list(range(hparams.train_size))),
                         drop_last=True)
    
    def val_dataloader(self):
        return DataLoader(PretrainingDatasetWrapper(stl10_unlabeled,
                                            debug=getattr(self.hparams, "debug", False)),
                          batch_size=self.hparams.batch_size, 
                          shuffle=False,
                          num_workers=cpu_count(),
                          sampler=SequentialSampler(list(range(hparams.train_size + 1, hparams.train_size + hparams.validation_size))),
                         drop_last=True)
    
    def forward(self, X):
        return self.model(X)
    
    def step(self, batch, step_name = "train"):
        (X, Y), y = batch
        embX, projectionX = self.forward(X)
        embY, projectionY = self.forward(Y)
        loss = self.loss(projectionX, projectionY)
        loss_key = f"{step_name}_loss"
        tensorboard_logs = {loss_key: loss}

        return { ("loss" if step_name == "train" else loss_key): loss, 'log': tensorboard_logs,
                        "progress_bar": {loss_key: loss}}
    
    def training_step(self, batch, batch_idx):
        return self.step(batch, "train")
    
    def validation_step(self, batch, batch_idx):
        return self.step(batch, "val")
    
    def validation_end(self, outputs):
        if len(outputs) == 0:
            return {"val_loss": torch.tensor(0)}
        else:
            loss = torch.stack([x["val_loss"] for x in outputs]).mean()
            return {"val_loss": loss, "log": {"val_loss": loss}}

    def configure_optimizers(self):
        optimizer = RMSprop(self.model.parameters(), lr=self.hparams.lr)
        return [optimizer], []

Initial hyper-parameters. Batch size of 128 works fine with EfficientNet-B0 on GTX1070. Note that I’ve limited training dataset to first 10k images from STL10 for convenience of running this blogpost in the form of Jupyter Notebook / Google Colab.

Important! SimCLR greatly benefits from large batch sizes - it should be set to as high value as possible given the GPU / cluster limits.

python

from argparse import Namespace
hparams = Namespace(
    lr=1e-3,
    epochs=50,
    batch_size=160,
    train_size=10000,
    validation_size=1000
)

Finding good initial-learning rate using LRFinder algorithm

I use pytorch-lightning’s built-in LRFinder algorithm to find initial learning rate.

python

module = ImageEmbeddingModule(hparams)
t = pl.Trainer(gpus=1)
lr_finder = t.lr_find(module)
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
CUDA_VISIBLE_DEVICES: [0]

   | Name  | Type            | Params
------------------------------------------
0 | model | ImageEmbedding  | 6 M  
1 | loss  | ContrastiveLoss | 0

python

lr_finder.plot(show=False, suggest=True)
lr_finder.suggestion()
0.000630957344480193

I also use W&B Logging to keep track of my experiments:

python

from pytorch_lightning.loggers import WandbLogger

python

hparams = Namespace(
    lr=0.000630957344480193,
    epochs=10,
    batch_size=160,
    train_size=20000,
    validation_size=1000
)
module = ImageEmbeddingModule(hparams)
logger = WandbLogger(project="simclr-blogpost")
logger.watch(module, log="all", log_freq=50)
trainer = pl.Trainer(gpus=1, logger=logger)

python

trainer.fit(module)
| Name  | Type            | Params
------------------------------------------
0 | model | ImageEmbedding  | 6 M   
1 | loss  | ContrastiveLoss | 0

After the training completes, embeddings are ready for usage for the downstream tasks.

Image classifier on top of SimCLR embeddings

Once the embeddings are trained, they can be used to train the classifier on top of them - either by fine tuning the whole network or by freezing the base network with embeddings and learning linear classifiers on top of it - I show the latter below.

Save the weights of neural network with embeddings

I save the whole network in a form on a checkpoint. Only the internal part of the network will be used later with the classifier (projection layers will be discarded).

python

checkpoint_file = "efficientnet-b0-stl10-embeddings.ckpt"

python

trainer.save_checkpoint(checkpoint_file)

python

trainer.logger.experiment.log_artifact(checkpoint_file, type="model")

Classifier module

Again, I define a custom module - this time it uses already existing embeddings and freezes the base model weights on demand. Note the SimCLRClassifier.embeddings are only the EfficientNet part of the whole network used before - the projection head is discarded.

python

class SimCLRClassifier(nn.Module):
    def __init__(self, n_classes, freeze_base, embeddings_model_path, hidden_size=512):
        super().__init__()
        
        base_model = ImageEmbeddingModule.load_from_checkpoint(embeddings_model_path).model
        
        self.embeddings = base_model.embedding
        
        if freeze_base:
            print("Freezing embeddings")
            for param in self.embeddings.parameters():
                param.requires_grad = False
                
        # Only linear projection on top of the embeddings should be enough
        self.classifier = nn.Linear(in_features=base_model.projection[0].in_features, 
                      out_features=n_classes if n_classes > 2 else 1)

    
    def forward(self, X, *args):
        emb = self.embeddings(X)
        return self.classifier(emb)

Classifier training code

Classifier training code again uses PyTorch lightning, so I’m skipping in-depth explanation.

python

from torch import nn
from torch.optim.lr_scheduler import CosineAnnealingLR

python

class SimCLRClassifierModule(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()
        hparams = Namespace(**hparams) if isinstance(hparams, dict) else hparams
        self.hparams = hparams
        self.model = SimCLRClassifier(hparams.n_classes, hparams.freeze_base, 
                                      hparams.embeddings_path,
                                      self.hparams.hidden_size)
        self.loss = nn.CrossEntropyLoss()
    
    def total_steps(self):
        return len(self.train_dataloader()) // self.hparams.epochs
    
    def preprocessing(seff):
        return transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
    
    def get_dataloader(self, split):
        return DataLoader(STL10(".", split=split, transform=self.preprocessing()),
                          batch_size=self.hparams.batch_size, 
                          shuffle=split=="train",
                          num_workers=cpu_count(),
                         drop_last=False)
    
    def train_dataloader(self):
        return self.get_dataloader("train")
    
    def val_dataloader(self):
        return self.get_dataloader("test")
    
    def forward(self, X):
        return self.model(X)
    
    def step(self, batch, step_name = "train"):
        X, y = batch
        y_out = self.forward(X)
        loss = self.loss(y_out, y)
        loss_key = f"{step_name}_loss"
        tensorboard_logs = {loss_key: loss}

        return { ("loss" if step_name == "train" else loss_key): loss, 'log': tensorboard_logs,
                        "progress_bar": {loss_key: loss}}
    
    def training_step(self, batch, batch_idx):
        return self.step(batch, "train")
    
    def validation_step(self, batch, batch_idx):
        return self.step(batch, "val")
    
    def test_step(self, batch, batch_idx):
        return self.step(Batch, "test")
    
    def validation_end(self, outputs):
        if len(outputs) == 0:
            return {"val_loss": torch.tensor(0)}
        else:
            loss = torch.stack([x["val_loss"] for x in outputs]).mean()
            return {"val_loss": loss, "log": {"val_loss": loss}}

    def configure_optimizers(self):
        optimizer = RMSprop(self.model.parameters(), lr=self.hparams.lr)
        schedulers = [
            CosineAnnealingLR(optimizer, self.hparams.epochs)
        ] if self.hparams.epochs > 1 else []
        return [optimizer], schedulers

It’s worth mentioning here that training with a frozen base model gives a great performance boost during training as the gradients need to be calculated only for a single layer. Additionally, by utilizing good embeddings, only a few epochs are required to reach a good quality classifier with single linear projection.

python

hparams_cls = Namespace(
    lr=1e-3,
    epochs=5,
    batch_size=160,
    n_classes=10,
    freeze_base=True,
    embeddings_path="./efficientnet-b0-stl10-embeddings.ckpt",
    hidden_size=512
)
module = SimCLRClassifierModule(hparams_cls)
logger = WandbLogger(project="simclr-blogpost-classifier")
logger.watch(module, log="all", log_freq=10)
trainer = pl.Trainer(gpus=1, max_epochs=hparams_cls.epochs, logger=logger)

python

lr_find_cls = trainer.lr_find(module)
| Name  | Type             | Params
-------------------------------------------
0 | model | SimCLRClassifier | 4 M   
1 | loss  | CrossEntropyLoss | 0  
LR finder stopped early due to diverging loss.

python

lr_find_cls.plot(show=False, suggest=True)
lr_find_cls.suggestion()
0.003981071705534969

python

hparams_cls = Namespace(
    lr=0.003981071705534969,
    epochs=5,
    batch_size=160,
    n_classes=10,
    freeze_base=True,
    embeddings_path="./efficientnet-b0-stl10-embeddings.ckpt",
    hidden_size=512
)
module = SimCLRClassifierModule(hparams_cls)

python

trainer.fit(module)
| Name  | Type             | Params
-------------------------------------------
0 | model | SimCLRClassifier | 4 M   
1 | loss  | CrossEntropyLoss | 0

Evaluation

Here I define a utility function for evaluating the model using the provided data loader. Note the transfer between GPU and CPU and storing all results in-memory might be not effective for large datasets!

python

from sklearn.metrics import classification_report
def evaluate(data_loader, module):
    with torch.no_grad():
        progress = ["/", "-", "\\", "|", "/", "-", "\\", "|"]
        module.eval().cuda()
        true_y, pred_y = [], []
        for i, batch_ in enumerate(data_loader):
            X, y = batch_
            print(progress[i % len(progress)], end="\r")
            y_pred = torch.argmax(module(X.cuda()), dim=1)
            true_y.extend(y.cpu())
            pred_y.extend(y_pred.cpu())
        print(classification_report(true_y, pred_y, digits=3))
        return true_y, pred_y

python

_ = evaluate(module.val_dataloader(), module)

bash

                  precision    recall  f1-score   support
    
               0      0.856     0.864     0.860       800
               1      0.714     0.701     0.707       800
               2      0.903     0.919     0.911       800
               3      0.678     0.599     0.636       800
               4      0.665     0.746     0.703       800
               5      0.633     0.564     0.597       800
               6      0.729     0.781     0.754       800
               7      0.678     0.709     0.693       800
               8      0.868     0.910     0.888       800
               9      0.862     0.801     0.830       800
    
        accuracy                          0.759      8000
       macro avg      0.759     0.759     0.758      8000
    weighted avg      0.759     0.759     0.758      8000

Summary

I hope you find my explanation of the SimCLR framework useful. Feel free to ask anything, comment, play with the notebook and share this post with the ML community!

Additional links and resources

SimCLR paper - arxiv.org/pdf/2002.05709.pdf
EfficientNet paper - arxiv.org/abs/1905.11946
this notebook on GitHub - link
this notebook on Google Colab - link

Custom classifier on top of BERT-like Language Model - guide

Marcin Zabłocki — Mon, 23 Mar 2020 23:00:00 +0000

Last time I wrote about training the language models from scratch, you can find this post here. Now it’s time to take your pre-trained lamnguage model at put it into good use by fine-tuning it for real world problem, i.e text classification or sentiment analysis. In this post I will show how to take pre-trained language model and build custom classifier on top of it. As in the previous post - I cover all of the important parts in details and I point out possible mistakes.

Why custom model? There are two answers. First is that the fun in deep learning begins only when you can do something custom with your model. Experiments, experiments and more experiments! Second thing is that by implmenting some parts on your own, you gain better understaing of different parts of the modeling itself, but also the whole training/fine-tuning process.

TL;DR

This post covers:

taking existing pre-trained language model and understanding it’s output - here I use PolBERTa trained for Polish language
building custom classification head on top of the LM
using fast tokenizers to efficiently tokenize and pad input text as well as prepare attention masks
preparing reproducible training code with PyTorch Lightning
finding good starting learning rate for the model
validating the trained model on PolEmo 2.0 dataset (benchmark for Polish language sentiment analysis with 4 classes)

The whole post is available as a Colab Notebook, linked below in Additional links and resources section.

Libraries and preparation

In this post I will use pretty standard set of ML-related libraries, including: transformers & tokenizers from HuggingFace, PyTorch Lightning, pandas, scikit-learn and LR Finder.

python

!pip install transformers tokenizers pytorch-lightning

After running the cell below, you will probably have to restart the notebook (one time only).

python

!git clone https://github.com/davidtvs/pytorch-lr-finder.git && cd pytorch-lr-finder && python setup.py install

python

import torch
from torch import nn
from typing import List
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelWithLMHead, AdamW, get_linear_schedule_with_warmup
import logging
import os
from functools import lru_cache
from tokenizers import ByteLevelBPETokenizer
from tokenizers.processors import BertProcessing
import pytorch_lightning as pl
from torch.utils.data import DataLoader, Dataset
import pandas as pd
from argparse import Namespace
from sklearn.metrics import classification_report
torch.__version__

Understanding pre-trained language model input / outputs

In order to build your classifier on top of pre-trained language model you must first understand it outputs. Usually it requries reading the related paper or (if possible) reading the documentation / github code. Fortunately, it’s easy for Transformers library - as the models are documented and the return value is described well. For my case (RoBERTa) the docs are here: https://huggingface.co/transformers/model_doc/roberta.html#robertamodel.

We’re interested what RobertaForMaskedLM returns. It this case, output format is the following:

python

sequence_output, pooled_output, (hidden_states), (attentions)

Let’s take a look at it ourselves:

python

tokenizer = AutoTokenizer.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1")
model = AutoModelWithLMHead.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1")

We are actually interested only in base model, without the LM head which is:

python

base_model = model.roberta

Let’s see what does output look like by running the model forward pass with some input. We need to tokenize the input first:

python

text = "Mówiąc otwarcie jesteście reliktami. Ilu wiedźminów jeszcze żyje? Kilku? Góra kilkunastu. Nie stanowicie k<mask>. Poza tym, Zakon nie każe sobie płacić za ratowanie ludzi od złego."
enc = tokenizer.encode_plus(text)
enc.keys()

dict_keys(['input_ids', 'token_type_ids', 'attention_mask'])

Tokenizer returns input_ids which are the tokens of your input text. Then token_type_ids and attention_mask - first one is used when you use LM for next-sentence prediction task, second one is to tell the model which tokens are from your input text and which are just padding tokens.

For classification, you only need input_ids and attention_mask as token types are created on the fly by the internal model’s forward function (source).

Let’s run the model. The .unsqueeze(0) is just to simulate batch processing (in this case, it will be batch size of 1).

python

# sequence_output, pooled_output, (hidden_states), (attentions)
out = base_model(torch.tensor(enc["input_ids"]).unsqueeze(0), torch.tensor(enc["attention_mask"]).unsqueeze(0))
out[0].shape, out[1].shape

(torch.Size([1, 12, 768]), torch.Size([1, 768]))

We are interested only in first element of the tuple while building classifier as the second element is used for next-sentence prediction task and it is literally said in the Transformers docs not to use it:

This output is usually not a good summary of the semantic content of the input, you’re often better with averaging or pooling the sequence of hidden-states for the whole input sequence. [source]

Output shape description:

Keep in mind that for bigger batch sizes you have to pad to either longest text’s lenght or to model’s max supported length (which usually for BERT-like models is 512).

The second dimension contains internal representation of each of the input tokens, so for properly tokenized data, i.e:

python

t = "Mówiąc otwarcie jesteście reliktami."
enc = tokenizer.encode_plus(t)
token_representations = base_model(torch.tensor(enc["input_ids"]).unsqueeze(0))[0][0]
print(enc["input_ids"])
print(tokenizer.decode(enc["input_ids"]))
print(f"Length: {len(enc['input_ids'])}")
print(token_representations.shape)

[0, 355, 304, 1748, 71, 12029, 28398, 807, 2578, 19022, 18, 2]
<s> Mówiąc otwarcie jesteście reliktami.</s>
Length: 12
torch.Size([12, 768])

Representation for each token is as follows:

0 = <s> token is represented by token_representations[0]
355 = M token is represented by token_representations[1]

and so on…

Build custom classification head on top of LM

OK, we understand the model’s output and sizes, now we can build custom classification head on top of it. Have you found new awesome activation function? Or maybe you want to try-out different pooling methods to better fit your case? You’re free to experiment. In this example, I will try out custom Mish activation function on top of straightforward classification head inspired by RoBERTa’s original one (source) and DistilBERT’s custom initialization. Because why not!

python

# from https://github.com/digantamisra98/Mish/blob/b5f006660ac0b4c46e2c6958ad0301d7f9c59651/Mish/Torch/mish.py
@torch.jit.script
def mish(input):
    return input * torch.tanh(F.softplus(input))
  
class Mish(nn.Module):
    def forward(self, input):
        return mish(input)

python

class PolBERTaSentimentModel(nn.Module):
    def __init__(self, base_model, n_classes, base_model_output_size=768, dropout=0.05):
        super().__init__()
        self.base_model = base_model
        
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(base_model_output_size, base_model_output_size),
            Mish(),
            nn.Dropout(dropout),
            nn.Linear(base_model_output_size, n_classes)
        )
        
        for layer in self.classifier:
            if isinstance(layer, nn.Linear):
                layer.weight.data.normal_(mean=0.0, std=0.02)
                if layer.bias is not None:
                    layer.bias.data.zero_()

    def forward(self, input_, *args):
        X, attention_mask = input_
        hidden_states, _ = self.base_model(X, attention_mask=attention_mask)
        
        # here I use only representation of <s> token, but you can easily use more tokens,
        # maybe do some pooling / RNNs... go crazy here!
        return self.classifier(hidden_states[:, 0, :])

Hint

As a sanity check, I encourage you to run forward pass of the model on dummy data to validate the code. It also better to do this on CPU, as the stack traces of potential exceptions will be more meaningful. Mysterious CUDA erros very often are caused by simple Python-level errors.

python

classifier = PolBERTaSentimentModel(AutoModelWithLMHead.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1").roberta, 3)

python

X = torch.tensor(enc["input_ids"]).unsqueeze(0)
attn = torch.tensor(enc["attention_mask"]).unsqueeze(0)

python

classifier((X, attn))

tensor([[0.5643, 0.2237, 0.1142]], grad_fn=<AddmmBackward>)

The output is correct - we’ve run the model with batch size and expected 3-class output (for CrossEntropyLoss).

Preparing fast tokenizers to use with text dataset

When you prepare your dataset you have two options while implementing Dataset class:

return tokenized text and class in __getitem__ function
return just text and class in __getitem__ function and implement Collate function - this function will be used by PyTorch DataLoader to prepare batches

Here I’ve chosen the second approach.

Default tokenizer loaded above (as for Transformers v2.5.1) uses Python implementation. In order to leverage full potential of parallel Rust tokenizers, we need to save the tokenizer’s internal data and then create instance of fast tokenizer with it.

python

!mkdir -p tokenizer

python

tokenizer.save_pretrained("tokenizer")

('tokenizer/vocab.json',
 'tokenizer/merges.txt',
 'tokenizer/special_tokens_map.json',
 'tokenizer/added_tokens.json')

python

!ls tokenizer

merges.txt  special_tokens_map.json  tokenizer_config.json  vocab.json

Once saved, implement CollateFn using fast tokenizers. As you can see, the whole batch can be tokenized in parallel and the tokenizer will also automatically return attention mask and pad input to max length of 512.

Important!

Ensure that all of the special tokens are in appropriate places and their IDs come from the actual tokenizer you’re using - they should not be hardcoded, as different models will have different IDs for padding, BOS, EOS and so on.

python

class TokenizersCollateFn:
    def __init__(self, max_tokens=512):
        t = ByteLevelBPETokenizer(
            "tokenizer/vocab.json",
            "tokenizer/merges.txt"
        )
        t._tokenizer.post_processor = BertProcessing(
            ("</s>", t.token_to_id("</s>")),
            ("<s>", t.token_to_id("<s>")),
        )
        t.enable_truncation(max_tokens)
        t.enable_padding(max_length=max_tokens, pad_id=t.token_to_id("<pad>"))
        self.tokenizer = t

    def __call__(self, batch):
        encoded = self.tokenizer.encode_batch([x[0] for x in batch])
        sequences_padded = torch.tensor([enc.ids for enc in encoded])
        attention_masks_padded = torch.tensor([enc.attention_mask for enc in encoded])
        labels = torch.tensor([x[1] for x in batch])
        
        return (sequences_padded, attention_masks_padded), labels

Getting data for classification

Here I use PolEmo 2.0 dataset, which is a benchmark for sentiment analysis task in Polish language.

python

!wget -O dataset_conll.zip https://clarin-pl.eu/dspace/bitstream/handle/11321/710/dataset_conll.zip?sequence=1&isAllowed=y

python

!unzip -n dataset_conll.zip

Archive:  dataset_conll.zip

python

train_path, val_path, test_path = [f"dataset_conll/all.text.{d}.txt" for d in ("train", "dev", "test")]
train_path, val_path, test_path

('dataset_conll/all.text.train.txt',
 'dataset_conll/all.text.dev.txt',
 'dataset_conll/all.text.test.txt')

This dataset has CSV-like format with __label__ text being the separator. Don’t ask. Let’s take a quick look at the data:

python

df = pd.read_csv(train_path, sep="__label__", header=None, names=["text", "label"], engine="python")
df.head()

	text	label
0	Na samym wejściu hotel śmierdzi . W pokojach j...	meta_minus_m
1	Ze swoimi dziećmi ( 10 lat i 2 latka ) chodzę ...	meta_plus_m
2	Chciał em zestaw do podstawowych prac coś przy...	meta_minus_m
3	Drogie Mamy . Mam o tym laktatorze takie same ...	meta_minus_m
4	własciwie nic się nie dowidział em , podczas b...	meta_minus_m

python

df.label.unique()

array(['meta_minus_m', 'meta_plus_m', 'meta_zero', 'meta_amb'],
      dtype=object)

You will need mapping between text label and integer:

python

label2int = {
  "meta_minus_m": 0,
  "meta_plus_m": 1,
  "meta_zero": 2,
  "meta_amb": 3
}

As you can see below, the Dataset itself is really simple as it returns only the text and label. The rest is handled by the TokenizersCollateFn above.

python

class PolEmo2Dataset(Dataset):
    def __init__(self, path):
        super().__init__()
        self.data_column = "text"
        self.class_column = "class"
        self.data = pd.read_csv(path, sep="__label__", header=None, names=[self.data_column, self.class_column],
                               engine="python")

    def __getitem__(self, idx):
        return self.data.loc[idx, self.data_column], label2int[self.data.loc[idx, self.class_column]]

    def __len__(self):
        return self.data.shape[0]

Sanity check:

python

ds = PolEmo2Dataset(train_path)
ds[13]

('Była m przekonana , że w Sopocie jest jeden hotel , który może zapewnić najlepszą jakość , tym większe było moje zaskoczenie krótkim pobytem służbowym w Małej Anglii . Miała m okazję mieszkać w pięknym i wygodnym pokoju o bardzo wysokim standardzie . Przemiła obsługa , co jest kluczowe i często niesieciowe hotele nie potrafią sobie z tym poradzić , tutaj nie było takiego problemu . Bardzo przyjemna strefa wellness oraz spa . Doskonała restauracja śniadaniowa , przepięknie urządzona , śniadanie bez zarzutu i kawa . . doskonała , każdy kawosz zrozumie . Brak widoku na morze został zupełnie zrekompensowany . Polecam i mam nadzieję , że taki poziom już pozostanie . Dziękuję za miłe wspomnienia . ',
 1)

Preparing reproducible training code with PyTorch Lightning

Most recently, PyTorch Lightning is my go to library for quick prototyping with PyTorch. You have clear API that is actually extension of the original PyTorch nn.Module one with all of the repeatable parts like training loop, validation loop, using GPUs, learning rate schedulers, gradient accumulation, tensorboard, checkpointing and many others handled for you.

There is only a thin API you have to implement, which is pl.LightningModule. Required parts are:

__init__ to initialize custom variables in your model
forward same as you would do in nn.Module - here I wrap my model inside of the pl.LightningModule but you could as well implement it directly here
training_step - take care of the forward pass for the model and return loss
validation_step and test_step - same as above but for the different learning phases
train_dataloader - DataLoader for your training dataset
val_dataloader and test_dataloader - same as above but for the different learning phases
configure_optimizers - prepare optimizer and (optionally) LR scheudles

python

class TrainingModule(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()
        self.model = PolBERTaSentimentModel(AutoModelWithLMHead.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1").roberta, 4)
        self.loss = nn.CrossEntropyLoss()
        self.hparams = hparams

    def step(self, batch, step_name="train"):
        X, y = batch
        loss = self.loss(self.forward(X), y)
        loss_key = f"{step_name}_loss"
        tensorboard_logs = {loss_key: loss}

        return { ("loss" if step_name == "train" else loss_key): loss, 'log': tensorboard_logs,
               "progress_bar": {loss_key: loss}}

    def forward(self, X, *args):
        return self.model(X, *args)

    def training_step(self, batch, batch_idx):
        return self.step(batch, "train")
    
    def validation_step(self, batch, batch_idx):
        return self.step(batch, "val")

    def validation_end(self, outputs: List[dict]):
        loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        return {"val_loss": loss}
        
    def test_step(self, batch, batch_idx):
        return self.step(batch, "test")
    
    def train_dataloader(self):
        return self.create_data_loader(self.hparams.train_path, shuffle=True)

    def val_dataloader(self):
        return self.create_data_loader(self.hparams.val_path)

    def test_dataloader(self):
        return self.create_data_loader(self.hparams.test_path)
                
    def create_data_loader(self, ds_path: str, shuffle=False):
        return DataLoader(
                    PolEmo2Dataset(ds_path),
                    batch_size=self.hparams.batch_size,
                    shuffle=shuffle,
                    collate_fn=TokenizersCollateFn()
        )
        
    @lru_cache()
    def total_steps(self):
        return len(self.train_dataloader()) // self.hparams.accumulate_grad_batches * self.hparams.epochs

    def configure_optimizers(self):
        optimizer = AdamW(self.model.parameters(), lr=self.hparams.lr)
        lr_scheduler = get_linear_schedule_with_warmup(
                    optimizer,
                    num_warmup_steps=self.hparams.warmup_steps,
                    num_training_steps=self.total_steps(),
        )
        return [optimizer], [{"scheduler": lr_scheduler, "interval": "step"}]

Hints and comments

hparams from __init__ is a Namespace-type object, which holds all of the models hyperparameters; PyTorch Lightning takes good care of them when checkpointing the model
*_step functions can be implemented with the same method with argument
*_dataloader functions can also be implemented in this way
in configure_optimizers the line: [{"scheduler": lr_scheduler, "interval": "step"}] takes care of the step-based LR scheduler - it does LR warmup and then linear decay, implementation comes from Transformers library

The code above is fairly simple and we don’t need anything more here. For more robust examples I send you to PyTorch Lightning docs.

Finding good starting learning rate for the model

There is a method of approximating “good” learning rate for the training. I first encountered it in Fast.ai Deep Learning course. This approach is called LR-Finder and boils down to plotting loss of the training batch (starting from really small LR) while increasing learning rate up to the poit where the loss blows up. Then, the “good” LR is one of the values that appears when the Loss/Learning Rate curve is decreasing.

Important!

Keep in mind that this method does not guarantee to provide good learning rate for every dataset. It’s just a method discovered to often work in practice.

Here I use implementation for PyTorch from here (installed at the top of this notebook).

You need to specify initial learning rate and number of steps. Other parameters are optional.

python

from torch_lr_finder import LRFinder
hparams_tmp = Namespace(
    train_path=train_path,
    val_path=val_path,
    test_path=test_path,
    batch_size=8,
    warmup_steps=100,
    epochs=3,
    lr=lr,
    accumulate_grad_batches=1,
)
module = TrainingModule(hparams_tmp)
criterion = nn.CrossEntropyLoss()
optimizer = AdamW(module.parameters(), lr=5e-7)
lr_finder = LRFinder(module, optimizer, criterion, device="cuda")
lr_finder.range_test(module.train_dataloader(), end_lr=100, num_iter=100, accumulation_steps=hparams_tmp.accumulate_grad_batches)
lr_finder.plot()
lr_finder.reset()

Stopping early, the loss has diverged
Learning rate search finished. See the graph with {finder_name}.plot()

From the plot above we can guess that something between 1e-5 and 1e-4 would be a good learning rate, as everyhing higher results in increased loss.

To ensure the right spot was chosen, you can plot it:

python

lr = 4e-5
lr

4e-05

python

lr_finder.plot(show_lr=lr)

Training the sentiment analysis classifier on PolEmo2.0

We know the learning rate, we have the code for the model. Now it’s time to run it.

python

hparams = Namespace(
    train_path=train_path,
    val_path=val_path,
    test_path=test_path,
    batch_size=8,
    warmup_steps=100,
    epochs=3,
    lr=lr,
    accumulate_grad_batches=1
)
module = TrainingModule(hparams)

python

# just to get rid of any leftovers before the training - useful when you run different cells multiple times
import gc; gc.collect()
torch.cuda.empty_cache()

Using Trainer from PyTorch Lightning

The pl.Trainer provides all-in-one solution for training, validation and testing. There is a lot of useful parameters for almost every aspect of training you can think of, everything is well documented here.

python

trainer = pl.Trainer(gpus=1, max_epochs=hparams.epochs, progress_bar_refresh_rate=10,
                     accumulate_grad_batches=hparams.accumulate_grad_batches)

python

trainer.fit(module)

Evaluating the trained model

As this dataset is small, I’ve used scikit-learn to get the overall classification report from the model.

Important!

Be aware that for big datasets this approach might not be optimal as it transfers the data from GPU to CPU for the whole dataset and puts every prediction in memory. For big datasets it’s better to implement metrics inside of the pl.LightningModule when the dataset is processed in batches.

Ensure that you disable gradients and put model into evaluation state (module.eval()) to disable dropouts during prediction.

python

with torch.no_grad():
    progress = ["/", "-", "\\", "|", "/", "-", "\\", "|"]
    module.eval()
    true_y, pred_y = [], []
    for i, batch_ in enumerate(module.test_dataloader()):
        (X, attn), y = batch_
        batch = (X.cuda(), attn.cuda())
        print(progress[i % len(progress)], end="\r")
        y_pred = torch.argmax(module(batch), dim=1)
        true_y.extend(y.cpu())
        pred_y.extend(y_pred.cpu())
print("\n" + "_" * 80)
print(classification_report(true_y, pred_y, target_names=label2int.keys(), digits=6))

python

    ________________________________________________________________________________
                  precision    recall  f1-score   support
    
    meta_minus_m   0.880878  0.828909  0.854103       339
     meta_plus_m   0.867521  0.894273  0.880694       227
       meta_zero   0.957983  0.966102  0.962025       118
        meta_amb   0.574324  0.625000  0.598592       136
    
        accuracy                       0.832927       820
       macro avg   0.820177  0.828571  0.823854       820
    weighted avg   0.837433  0.832927  0.834617       820

Additional links and resources

this notebook in Google Colab - https://colab.research.google.com/drive/1sajgpLTrTJDzRSlxycy8aE6ysxGqaT18
this notebook on GitHub - https://github.com/marrrcin/transformers-sentiment-analysis
PyTorch Lightning documentation - https://pytorch-lightning.readthedocs.io/en/0.7.1/

Training RoBERTa from scratch - the missing guide

Marcin Zabłocki — Sun, 15 Mar 2020 23:00:00 +0000

After hours of research and attempts to understand all of the necessary parts required for one to train custom BERT-like model from scratch using HuggingFace’s Transformers library I came to conclusion that existing blog posts and notebooks are always really vague and do not cover important parts or just skip them like they weren’t there - I will give a few examples, just follow the post.

I’ve decided to get my hands dirty and try to train transformer language model (BERT or other) for Polish language. Here you will find all of the missing pieces that I have encountered while doing it.

Many thanks to Egnyte Inc. for providing me with necessary resources to run the training on powerful machine with 4x NVIDIA P100 GPUs.

TL;DR

This post covers:

what mistakes to avoid
how to prepare own dataset with details - I use Polish Wikipedia and Wolne Lektury datasets
how to prepare training configuration
how to train new tokenizers using HugginFace’s Rust tokenizers
how to use run_language_modeling.py script from HuggingFace library to train your model - what parameters to set, how to deal with the script
how to monitor training process
how to share the model on HuggingFace’s model hub

Getting the dataset

The most important part when dealing with language models is to have solid dataset with text in the language you will be modeling. In my proof of concept I’ve started with Polish Wikipedia and then extended my dataset with Wolne Lektury (repository of public domain polish books and poems).

Preparing Polish Wikipedia dump

So there is a “just get wikipedia” part that noone covers in details but I find it important. First of all, on the wikipedia dumps page there are many files that you can download. I’ve decided to download dump from 2020-02-20 labeled as:

Articles, templates, media/file descriptions, and primary meta-pages, in multiple bz2 streams, 100 pages per stream

which consisted of 7 files with names like: plwiki-20200220-pages-articles-multistream*.xml*.bz2. Why? That kind of split allows to simpify preprocessing - I don’t have to have lots of RAM to processs smaller files. Another benefit is that I can just hold-out one or two of those files for validation phase after my language model learning is complete.

Every file is a huge XML containing articles with mediawiki-specific markup language. I used this tool to extract well formated JSONs from them: https://github.com/attardi/wikiextractor (clone the repo).

So now I have a list of files for both training and evaluation:

bash

    # train.txt
    plwiki-20200220-pages-articles-multistream1.xml-p1p169750
    plwiki-20200220-pages-articles-multistream2.xml-p169751p510662
    plwiki-20200220-pages-articles-multistream3.xml-p510663p1056310
    plwiki-20200220-pages-articles-multistream4.xml-p1056311p1831508
    plwiki-20200220-pages-articles-multistream5.xml-p1831509p3070393
    plwiki-20200220-pages-articles-multistream6.xml-p4570394p4720470

and

bash

    # eval.txt
    plwiki-20200220-pages-articles-multistream6.xml-p3070394p4570393

To extract them just run:

bash

cat train.txt | xargs -I@  python wikiextractor/WikiExtractor.py @ --bytes=100M --json --output="./plwiki-json/train/@"

After extracting you will get the following structure:

bash

├── eval
│   └── plwiki-20200220-pages-articles-multistream6.xml-p3070394p4570393
│       └── AA
│           ├── wiki_00
│           ├── wiki_01
│           ├── wiki_02
│           └── wiki_03
└── train
    ├── plwiki-20200220-pages-articles-multistream1.xml-p1p169750
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       └── wiki_02
    ├── plwiki-20200220-pages-articles-multistream2.xml-p169751p510662
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       └── wiki_02
    ├── plwiki-20200220-pages-articles-multistream3.xml-p510663p1056310
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       └── wiki_02
    ├── plwiki-20200220-pages-articles-multistream4.xml-p1056311p1831508
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       ├── wiki_02
    │       └── wiki_03
    ├── plwiki-20200220-pages-articles-multistream5.xml-p1831509p3070393
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       ├── wiki_02
    │       └── wiki_03
    └── plwiki-20200220-pages-articles-multistream6.xml-p4570394p4720470
        └── AA
            └── wiki_00

Extracting text from wikipedia

After you get JSON lines from wikipedia dump, you need to transform articles in JSON text into plaintext consumable by the language model training script.

Important!

I’ve probably made a mistake here which is important later on! I first tokenized the text by sentences and outputted each sentence in a separate line because of the ambiguity of available resources about dataset preparation. I think that it affects the quality of both tokenization and training. Do not make that mistake.

After digging throughout the HuggingFace’s github I’ve finally got an answer that there should be one document per line. OK - but what if the text is really long? We will be dealing not only with long wiki articles but later on - with books that have thousands of characters and current SOTA language models have limit of around 512-768 tokens (depends on the model). The possible solution for that is to use sliding window over sentences in the input text. For more details see the following thread: https://github.com/huggingface/transformers/issues/2693

To extract the text and tokenize the input I’ve used nltk with custom polish abbreviations - thanks to Krzysztof Sopyła’s awesome gist: https://gist.github.com/ksopyla/f05fe2f48bbc9de895368b8a7863b5c3.

You can find my extraction notebook here: https://gist.github.com/marrrcin/e383b75a5d0dad42048847d97965e037. I decided to skip it here to not pollute the post.

Key takeaways from preprocessing:

tokenize article into sentences
use sliding window over sentences to handle long texts - I used sliding window of size 4

Once the extraction completes, you will have training file like this:

bash

Sumatra Północna Sumatra Północna (indonez. "Sumatera Utara") – prowincja w Indonezji w północnej części Sumatry. Obejmuje również wyspę Nias i wyspy Batu leżące na Oceanie Indyjskim. Powierzchnia 70 787 km²; 12 391 tys. mieszkańców (2005); stolica Medan.
"Sumatera Utara") – prowincja w Indonezji w północnej części Sumatry. Obejmuje również wyspę Nias i wyspy Batu leżące na Oceanie Indyjskim. Powierzchnia 70 787 km²; 12 391 tys. mieszkańców (2005); stolica Medan. Większą część powierzchni zajmują góry Barisan (Sinabung 2460 m n.p.m.) z licznymi wulkanami i jeziorem Toba.
Obejmuje również wyspę Nias i wyspy Batu leżące na Oceanie Indyjskim. Powierzchnia 70 787 km²; 12 391 tys. mieszkańców (2005); stolica Medan. Większą część powierzchni zajmują góry Barisan (Sinabung 2460 m n.p.m.) z licznymi wulkanami i jeziorem Toba. Bliżej wybrzeży nizinne tereny w dużym stopniu pokryte bagnami.
Powierzchnia 70 787 km²; 12 391 tys. mieszkańców (2005); stolica Medan. Większą część powierzchni zajmują góry Barisan (Sinabung 2460 m n.p.m.) z licznymi wulkanami i jeziorem Toba. Bliżej wybrzeży nizinne tereny w dużym stopniu pokryte bagnami. Ludność prowincji tworzą Malajowie wyznający islam, Batakowie będący w większości chrześcijanami (gł. protestantami) oraz przybysze z Jawy, Chin, Indii.
Większą część powierzchni zajmują góry Barisan (Sinabung 2460 m n.p.m.) z licznymi wulkanami i jeziorem Toba. Bliżej wybrzeży nizinne tereny w dużym stopniu pokryte bagnami. Ludność prowincji tworzą Malajowie wyznający islam, Batakowie będący w większości chrześcijanami (gł. protestantami) oraz przybysze z Jawy, Chin, Indii. Główne miasta: Medan, Binjai, Pematang Siantar, Tebing Tinggi

Prepare Wolne Lektury dump

For extraction of polish books I have crawled https://wolnelektury.pl/katalog/ site. It’s pretty straightforward process - all of the books can be downloaded in well formatted TXT files. The preprocessing steps are exactly the same as for wiki, namely: tokenize book into sentences, use sliding window to transform long books into smaller chunks without dropping any data - this time I used sliding window of 8 sentences. Gist of preprocessing notebook: https://gist.github.com/marrrcin/bcc115fbadf79eba9d9c8ca711da9e20.

Training new tokenizer

At this point, I’ve decided to go with RoBERTa model. Model you choose determines the tokenizer that you will have to train. For RoBERTa it’s a ByteLevelBPETokenizer, for BERT it would be BertWordPieceTokenizer (both from tokenizers library). Training the tokenizer is super fast thanks to the Rust implementation that guys at HuggingFace have prepared (great job!).

python


from pathlib import Path

from tokenizers import ByteLevelBPETokenizer
from glob import glob
paths = list(
    glob("your_input_dataset_files/*.txt")
)
# Initialize a tokenizer
tokenizer = ByteLevelBPETokenizer(lowercase=False)

# Customize training
tokenizer.train(files=paths, vocab_size=32000, min_frequency=3, special_tokens=[
    "<s>",
    "<pad>",
    "</s>",
    "<unk>",
    "<mask>",
])

# Save files to disk
import os
OUT_DIR = "polish_tokenizer_bpe_32k"
os.makedirs(OUT_DIR, exist_ok=True)
tokenizer.save(OUT_DIR, "pl")

Here, the special_tokens part is super important - those tokens need to be exactly in the same order, because otherwise the model will crash!.

For BERT tokenizer however, you will have:

python

from tokenizers import BertWordPieceTokenizer, ByteLevelBPETokenizer
# Initialize a tokenizer
tokenizer = BertWordPieceTokenizer(lowercase=False, handle_chinese_chars=False)

Adjusting the `run_language_modeling.py` script

The version I used and modified comes from v.2.5.1 release of HuggingFace transformers: https://github.com/huggingface/transformers/blob/v2.5.1/examples/run_language_modeling.py

Dataset - check, tokenizer - check. Now it’s time for the model training. Initially, I wanted to use the run_language_modeling.py script from HuggingFace git repo without changes, but it’s it does not support fast Rust-based tokenizers (as for 2020-03-16). Here I mention the minimal changes I made to make it work.

Fast tokenizers support in `run_language_modeling.py`

In order to plug fast tokenizer I’ve trained above into this script I had to modify the LineByLineTextDataset that’s provided there. The final version looks like this:

python

class LineByLineTextDataset(Dataset):
    def __init__(self, t: PreTrainedTokenizer, args, file_path: str, block_size=512):
        assert os.path.isfile(file_path)
        logger.info("Creating features from dataset file at %s", file_path)
        
        # -------------------------- CHANGES START
        bert_tokenizer = os.path.join(args.tokenizer_name, "vocab.txt")
        if os.path.exists(bert_tokenizer):
            logger.info("Loading BERT tokenizer")
            from tokenizers import BertWordPieceTokenizer
            tokenizer = BertWordPieceTokenizer(os.path.join(args.tokenizer_name, "vocab.txt"), handle_chinese_chars=False, lowercase=False)
            tokenizer.enable_truncation(512)
        else:
            from tokenizers import ByteLevelBPETokenizer
            from tokenizers.processors import BertProcessing
            logger.info("Loading RoBERTa tokenizer")
            
            tokenizer = ByteLevelBPETokenizer(
                os.path.join(args.tokenizer_name, "vocab.json"),
                os.path.join(args.tokenizer_name, "merges.txt")
            )
            tokenizer._tokenizer.post_processor = BertProcessing(
                ("</s>", tokenizer.token_to_id("</s>")),
                ("<s>", tokenizer.token_to_id("<s>")),
            )
            tokenizer.enable_truncation(max_length=512)

        logger.info("Reading file %s", file_path)
        with open(file_path, encoding="utf-8") as f:
            lines = [line for line in f.read().splitlines() if (len(line) > 0 and not line.isspace())]

        logger.info("Running tokenization")
        self.examples = tokenizer.encode_batch(lines)
        
        # -------------------------- CHANGES END

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        return torch.tensor(self.examples[i].ids, dtype=torch.long)

So the only changes are in the way of processing input args.tokenizer_name parameter - if the folder has vocab.txt it’s BERT tokenizer, RoBERTa otherwise. Thanks to the tokenizer.encode_batch(lines) the whole dataset is tokenized in parallel.

Important!

If you don’t have enough RAM to fit the whole dataset into memory, you have to implement your own dataset. My input dataset - Polish wikipedia with sliding window of 4 + Wolne Lektury with sliding window of 8 took ~235GB of RAM.

Preparing training configuration

This part is the most tricky one, as there are a lot of magic numbers in the posts / notebooks I’ve found, even in the official one https://huggingface.co/blog/how-to-train. They just skip it… Start with creating the directory for the configuration, say MyRoBERTaConfig. Copy the tokenizer files and rename them. Be sure to finally have:

bash

    config.json
    tokenizer_config.json
    merges.txt
    vocab.json

in this directory.

Training configuration for RoBERTa

You know what? You don’t have to specify any model-specific parameters if you want to go with the defaults! The only thing that you actually need to specify is the vocabulary size, because it determines the output layer size (Language Model is predicting the tokens after all…). For my RoBERTa I have the following config.json:

javascript

{"architectures": ["RobertaForMaskedLM"], "max_position_embeddings": 514, "vocab_size": 32000}

The architectures is self explanatory - this is the language model you will be training, vocab_size comes from the tokenizer you have created.

Important!

The tricky one is max_position_embeddings - no one will tell you that, but if you don’t set this here, the script will just crash. Awesome. I believe that for BERT model it’s not required.

Tokenizer configuration for RoBERTa is simple (tokenizer_config.json). However it’s probably not required, because we’ve overridden the LineByLineTextDataset dataset with tokenizer.

javascript

{"max_len": 512}

Training your own RoBERTa language model

Here I will show two types of run_language_modeling.py script execution - one for training from scratch and the second - for fine tuning.

Training from scratch

This is the usual case for any new language model. First, the script:

bash

#!/bin/bash
export TRAIN_FILE=plwiki.train.txt
export EVAL_FILE=plwiki.eval.txt
# run_language_modeling_with_tokenizers.py -- it's the version with support for fast tokenizers, see above
python run_language_modeling_with_tokenizers.py \
    --train_data_file $TRAIN_FILE \
    --eval_data_file $EVAL_FILE \
    --output_dir ./MyRoBERTa \
    --model_type roberta \
    --mlm \
    --config_name ./MyRoBERTaConfig \
    --tokenizer_name ./MyRoBERTaConfig \
    --do_train \
    --do_eval \
    --line_by_line \
    --learning_rate 1e-5 \
    --num_train_epochs 5 \
    --save_total_limit 20 \
    --save_steps 5000 \
    --per_gpu_train_batch_size 8 \
    --warmup_steps=10000 \
    --logging_steps=100 \
    --gradient_accumulation_steps=4 \
    --seed 666 --block_size=512

Arguments explanation: (only the non-trivial)

--output_dir - here you will get all of the model checkpoints, saved as often as specified in --save_steps argument
--mlm - configure the script to run Masked Language Model task. The script will fortunately yell at you if the --model_type is not appropriate for this task
--line_by_line - we say to use LineByLineTextDataset, the script will use my overridden version with tokenizers
--save_total_limit - how many checkpoint to keep on the disk - the oldest one will be removed once this limit is exceeded. It’s good to set it here in order to not clog your disk
--per_gpu_train_batch_size - setting this parameter might require a few failed runs fro you. If you don’t know how big batch size you can fit into your GPU, just sample your training data, to say 10k lines and launch the script a few times. Monitor the GPU usage by using watch -n 1 nvidia-smi command. Set it to as high number as possible. For Nvidia P100 and 512 --block_size, it’s 8.
--gradient_accumulation_steps - this parameter configures how many steps you will pass in forward pass of the network before running backward pass, it’s useful when you GPU cannot handle large batches
--warmup_steps - configures learning rate scheduler; as per BERT paper recommendation - it’s good to do warmup and then warmdown. Number of steps depends on your training dataset size
--logging_steps - tells how often to store Tensorboard logs - the script only outputs loss and learning rate

Training from checkpoint (finetuning)

After I’ve pretrained my model on wikipedia only, I decided to finetune it further on the Wolne Lektury dataset. Training script has produced a few snapshots, from which every one has the following structure:

bash

├── checkpoint-20000
│   ├── config.json
│   ├── merges.txt
│   ├── optimizer.pt
│   ├── pytorch_model.bin
│   ├── scheduler.pt
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   ├── training_args.bin
│   └── vocab.json

Important!

If you’re changing the input dataset, you cannot use existing scheduler.pt and optimizer.pt as you’re not “resuming” the training. You should remove (but backup first!) those files if you want to fine tune on another dataset.'

OK, now the script itself:

bash

#!/bin/bash
export TRAIN_FILE=wolne-lektury.sliding8.txt
export EVAL_FILE=plwiki.eval.txt

python run_language_modeling_with_tokenizers.py \
    --train_data_file $TRAIN_FILE \
    --eval_data_file $EVAL_FILE \
    --output_dir ./MyRoBERTa_with_another_dataset \
    --model_type roberta \
    --model_name_or_path ./MyRoBERTa/checkpoint-30000 \
    --mlm \
    --config_name ./MyRoBERTaConfig \
    --tokenizer_name ./MyRoBERTaConfig \
    --do_train \
    --do_eval \
    --line_by_line \
    --learning_rate 5e-5 \
    --num_train_epochs 3 \
    --save_total_limit 20 \
    --save_steps 5000 \
    --per_gpu_train_batch_size 8 \
    --warmup_steps=5000 \
    --logging_steps=100 \
    --gradient_accumulation_steps=4 \
    --mlm_probability=0.2 \
    --seed 666 --block_size=512

Arguments - notable mentions:

--mlm_probability=0.2 - this parameter controls the percentage of the tokens you mask during training; default is 0.15, I’ve decided to change it to make the training more difficult to the model
--model_name_or_path ./MyRoBERTa/checkpoint-30000 - this should point to the checkpoint of your previously trained model; remember the note above!

Training phase

I highly recommend you to launch the script using screen tool, which allows to easily spawn and manage multiple SSH sessions - you will be able to disconnect from the server while the training runs. And it will run for days… so brace yourself. You can use screen -UL command (-U for enabling UTF-8 mode for progress bars and -L to log session output to file - it’s helpful when something crashes).

During training, the script will output Tensorboard logs into runs directory. You can monitor them by running tensorboard command:

bash

tensorboard --logdir=runs

It will launch small webserver on localhost:6006 and you will be able to monitor the training.

Training summary for Polish RoBERTa a.k.a PolBERTa

I’ve run my training in three phases:

From scratch on Polish Wikipedia only (1 sentence per line) for 370k steps using learning rate 1e-5 with 10k warmup steps.
Fine tuning on Wolne Lektury only (8 sentences per line) for 60k steps, starting after checkpoint from 1. Learning rate 5e-5 with 5k warmup steps.
Fine tuning on Wolne Lektury (8 sentences per line) and Polish Wikipedia (4 sentences per line) for 30k steps. Learning rate 5-e5 with 10k warmup steps.

Tensorboard for all of the runs can be viewed here https://tensorboard.dev/experiment/aYrAE9uMTxGKRtLUOFjprg/

Exporting the model

If you want to share the model with the NLP community (which I highly encourage you to do!) you need to export it in appropriate format. Prepare the path for your model’s latest checkpoint and then run the following code:

python

from transformers import AutoModelWithLMHead, AutoTokenizer
import os
directory = "/path/to/your/model/checkpoint-30000"
model = AutoModelWithLMHead.from_pretrained(directory)
tokenizer = AutoTokenizer.from_pretrained(directory)
out = "MyBERTa-base-cased-v1"
os.makedirs(out, exist_ok=True)
model.save_pretrained(out)
tokenizer.save_pretrained(out)

Then use transformers-cli util to upload the model:

bash

transformers-cli upload ./MyBERTa-base-cased-v1/

Using my pre-trained model

I’ve uploaded my pre-trained RoBERTa to HuggingFace’s model hub: marrrcin/PolBERTa-base-polish-cased-v1 Usage:

python

from transformers import AutoModelWithLMHead, AutoTokenizer, pipeline

model = AutoModelWithLMHead.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1")
tokenizer = AutoTokenizer.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1")

fill_mask = pipeline('fill-mask', model=model, tokenizer=tokenizer)
fill_mask("Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad <mask>, a postępujesz tak jakbyś miał niecałe dziesięć.")

Output:

javascript

[{'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad sto, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.1013106107711792,
  'token': 1675},
 {'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad tysiąc, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.09420681744813919,
  'token': 31144},
 {'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad połowę, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.0891065001487732,
  'token': 13627},
 {'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad dwadzieścia, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.07204979658126831,
  'token': 18811},
 {'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad godzinę, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.059559762477874756,
  'token': 20775}]

Follow up and TODO

Here are a few follow up things and my planned future work on Polish RoBERTa:

prepare model card for HuggingFace's model hub
evaluate PolBERTa on PolEmo2.0 and other downstream tasks
re-train PolBERTa again, using bigger vocab size for tokenizer
fine tune model further on another open datasets

If you find any bug in the post or something is not clear - feel free to post a comment, I will be glad to improve this short practical guide.

Thanks to…

I would like to thank Egnyte Inc. for providing me with required resources to train the language model and also Darek Kłeczek (creator of original PolBERT) for tips related to training BERT.

Additional links and resources

PolBERT - original BERT-base-uncased LM trained by Darek Kłeczek
HuggingFaces original guide for training new language models

Using ResNet for MNIST in PyTorch

Marcin Zabłocki — Sun, 06 Jan 2019 23:00:00 +0000

During last year (2018) a lot of great stuff happened in the field of Deep Learning. One of those things was the release of PyTorch library in version 1.0. PyTorch is my personal favourite neural network/deep learning library, because it gives the programmer both high level of abstraction for quick prototyping as well as a lot of control when you want to dig deeper. Alongside that, PyTorch does not force you into learning any new API conventions, because everything that you define in PyTorch - from the network architecture, throught data loading to custom loss functions is defined in plain Python, using either ordinary functions or object oriented style.

In this post I will show you how to get started with PyTorch by explaining how to use pre-defined ResNet architecture to create image classifier for the MNIST dataset. I assume that you have some minimal prior knowledge of Neural Networks and Deep Learning.

This post is deprecated, see it's updated version:
https://zablo.net/blog/post/pytorch-resnet-mnist-jupyter-notebook-2021/

TL;DR

Quick start for PyTorch by showing how to adjust ResNet architecture to create deep learning image classifier for MNIST dataset with 99% accuracy/precision/recall after only 5 epochs.

Do not reinvent the wheel

When it comes to image classification using deep learning, there was already a lot of research done in the terms of finding the best network architecture. In the famous ImageNet competition, there were year-to-year improvements and competitors came up with new ideas every time. One of the architectures that became very succesfull in ImageNet competition was ResNet (2015). I will show you how to take existing network definition from torchvision library and tweak for use with MNIST dataset.

Digging into the ResNet

ResNet were originally designed for ImageNet competition, which was a color (3-channel) image classification task with 1000 classes. MNIST dataset howerver only contains 10 classes and it’s images are in the grayscale (1-channel). So there are two things to change in the original network.

I will take ResNet18 from torchvision library (official PyTorch module with network architectures, image transformations and others). Let’s see how it’s implemented there:

python

# from: https://github.com/pytorch/vision/blob/21153802a3086558e9385788956b0f2808b50e51/torchvision/models/resnet.py#L167

def resnet18(pretrained=False, **kwargs):
    """Constructs a ResNet-18 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
    return model

So they just create object of class torchvision.models.resnet.ResNet with appropriate parameters. Now, what happens when we train the network? In PyTorch, the forward function of network class is called - it represent forward pass of data through the network. ResNets forward looks like this:

python

# from: https://github.com/pytorch/vision/blob/21153802a3086558e9385788956b0f2808b50e51/torchvision/models/resnet.py#L149

def forward(self, x):
    x = self.conv1(x)
    x = self.bn1(x)
    ## ... skipped a few lines ...
    return x

So the first layer (input) is conv1. It’s defined in the ResNet constructor like this:

python

# from: https://github.com/pytorch/vision/blob/21153802a3086558e9385788956b0f2808b50e51/torchvision/models/resnet.py#L104

self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)

As you can see, ResNet takes 3-channel (RGB) image. Standard input image size for this network is 224x224px.

Adjusting ResNet for MNIST

Now that we know what’s inside the ResNet, let’s adjust it for our needs.

python

class MnistResNet(ResNet):
    def __init__(self):
        super(MnistResNet, self).__init__(BasicBlock, [2, 2, 2, 2], num_classes=10)
        self.conv1 = torch.nn.Conv2d(1, 64, 
            kernel_size=(7, 7), 
            stride=(2, 2), 
            padding=(3, 3), bias=False)

I’ve changed input layer, to take single-channel image and set the number of classes to 10.

And it’s done. We have MNIST-ready ResNet network architecture.

PyTorch training loop boilerplate

Unfortunately, there is no built-in training class in the plain PyTorch library, so here is a boilerplate code for training any network (you can copy and paste it). It’s more or less based on the training loop code from official PyTorch documentation enriched with precision/recall/F1/accuracy metrics calculation, progress bar and pretty-printing.

Imports:

python

from torchvision.models.resnet import ResNet, BasicBlock
from torchvision.datasets import MNIST
from tqdm.autonotebook import tqdm
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score
import inspect
import time
from torch import nn, optim
import torch
from torchvision.transforms import Compose, ToTensor, Normalize, Resize
from torch.utils.data import DataLoader

Training loop for any PyTorch model (with evaluation)

I’ve added as many neccessary comments as possible, to ease the reading and explain confusing bits.

python

# model:
model = YourModelHere()

# params you need to specify:
epochs = 5
train_loader, val_loader = # put your data loader here
loss_function = nn.CrossEntropyLoss() # your loss function, cross entropy works well for multi-class problems

# optimizer, I've used Adadelta, as it wokrs well without any magic numbers
optimizer = optim.Adadelta(model.parameters())

start_ts = time.time()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

losses = []
batches = len(train_loader)
val_batches = len(val_loader)

# loop for every epoch (training + evaluation)
for epoch in range(epochs):
    total_loss = 0

    # progress bar (works in Jupyter notebook too!)
    progress = tqdm(enumerate(train_loader), desc="Loss: ", total=batches)

    # ----------------- TRAINING  -------------------- 
    # set model to training
    model.train()
    
    for i, data in progress:
        X, y = data[0].to(device), data[1].to(device)
        
        # training step for single batch
        model.zero_grad()
        outputs = model(X)
        loss = loss_function(outputs, y)
        loss.backward()
        optimizer.step()

        # getting training quality data
        current_loss = loss.item()
        total_loss += current_loss

        # updating progress bar
        progress.set_description("Loss: {:.4f}".format(total_loss/(i+1)))
        
    # releasing unceseccary memory in GPU
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    # ----------------- VALIDATION  ----------------- 
    val_losses = 0
    precision, recall, f1, accuracy = [], [], [], []
    
    # set model to evaluating (testing)
    model.eval()
    with torch.no_grad():
        for i, data in enumerate(val_loader):
            X, y = data[0].to(device), data[1].to(device)

            outputs = model(X) # this get's the prediction from the network

            val_losses += loss_function(outputs, y)

            predicted_classes = torch.max(outputs, 1)[1] # get class from network's prediction
            
            # calculate P/R/F1/A metrics for batch
            for acc, metric in zip((precision, recall, f1, accuracy), 
                                   (precision_score, recall_score, f1_score, accuracy_score)):
                acc.append(
                    calculate_metric(metric, y.cpu(), predicted_classes.cpu())
                )
          
    print(f"Epoch {epoch+1}/{epochs}, training loss: {total_loss/batches}, validation loss: {val_losses/val_batches}")
    print_scores(precision, recall, f1, accuracy, val_batches)
    losses.append(total_loss/batches) # for plotting learning curve
print(f"Training time: {time.time()-start_ts}s")

+ helper functions:

python

def calculate_metric(metric_fn, true_y, pred_y):
    # multi class problems need to have averaging method
    if "average" in inspect.getfullargspec(metric_fn).args:
        return metric_fn(true_y, pred_y, average="macro")
    else:
        return metric_fn(true_y, pred_y)
    
def print_scores(p, r, f1, a, batch_size):
    # just an utility printing function
    for name, scores in zip(("precision", "recall", "F1", "accuracy"), (p, r, f1, a)):
        print(f"\t{name.rjust(14, ' ')}: {sum(scores)/batch_size:.4f}")

Loading MNIST dataset and training the ResNet

One last bit is to load the data. As ResNets in PyTorch take input of size 224x224px, I will rescale the images and also normalize the numbers. Normalization helps the network to converge (find the optimum) a lot faster.

Remember to normalize the data using parameters from training dataset only, as doing otherwise (i.e calculating mean and standard deviation on the whole dataset) will cause data leak between the datasets causing the results to be incorrect!

Dataset

python

def get_data_loaders(train_batch_size, val_batch_size):
    mnist = MNIST(download=False, train=True, root=".").train_data.float()
    
    data_transform = Compose([ Resize((224, 224)),ToTensor(), Normalize((mnist.mean()/255,), (mnist.std()/255,))])

    train_loader = DataLoader(MNIST(download=True, root=".", transform=data_transform, train=True),
                              batch_size=train_batch_size, shuffle=True)

    val_loader = DataLoader(MNIST(download=False, root=".", transform=data_transform, train=False),
                            batch_size=val_batch_size, shuffle=False)
    return train_loader, val_loader

Training & evaluation

I’ve run 5 epochs on Nvidia Tesla K80 GPU. Training took ~35 minutes. Final scores are around 0.99 for precision, recall and accuracy.

Summary

Whenever you encounter image classification task, it’s better to find existing, battle tested deep neural network architecture instead of spending the time trying to build your own. The time you save on the architecture design should be spent on polishing your dataset or finding more examples.

I hope this post will help you whenever you will be building your image classifier.

Additional links & resources

Jupyter notebook for this post: https://github.com/marrrcin/pytorch-resnet-mnist
good blog post explaining the ResNets: https://medium.com/@14prakash/understanding-and-implementing-architectures-of-resnet-and-resnext-for-state-of-the-art-image-cf51669e1624

pytorch on Marcin Zabłocki blog

Robust Subspace Recovery Layer implementation in PyTorch

TL;DR

Idea behind Robust Space Recovery Layer

What is the RSR Layer?

How to implement the RSR Layer in PyTorch?

RSR Auto Encoder in PyTorch

Training the autoencoder with RSR Layer in PyTorch (& PyTorch Lightning)

Corrupted MNIST dataset (with outliers)

RSRAutoencoder module in PyTorch Lightning

Analysing RSR Autoencoder outputs

Unsupervised anomaly detection using RSR Autoencoder

Summary

Additional links & resources

13 features of PyTorch that you should know

TL;DR

1. DatasetFolder

2. Stop using .to(device) that much: zeros_like / ones_like etc.

3. Register Buffer (a.k.a nn.Module.register_buffer)

4. Built-in Identity()

5. Pairwise distances: torch.cdist

6. Cosine similarity: F.cosine_similarity

Batched cosine similarity in PyTorch

7. Normalizing vectors: F.normalize

Batched L2 normalization in PyTorch

8. Linear layer + chunking trick (torch.chunk)

9. Masked select (torch.masked_select)

Applying mask on tensor directly

10. Conditional tensors with torch.where

11. Filling tensors with values at given positions (Tensor.scatter)

PyTorch scatter example on 2D data

12. Image interpolation within network (F.interpolate)

13. Making grid of images (torchvision.utils.make_grid)

Summary

Using ResNet for MNIST in PyTorch 1.7

TL;DR

Use ResNet from torchvision

Load the dataset

Do not write training loops in PyTorch

Inference with PyTorch Lightning

Summary

Additional links & resources

Understanding & implementing SimCLR in PyTorch - an ELI5 guide

TL;DR

Understanding SimCLR framework

Contrastive loss function

Theory behind contrastive loss function

Contrastive loss function - implementation in PyTorch, ELI5 version

Explanation

Contrastive loss function - implementation in PyTorch, vectorized version

Pre-training image embeddings using SimCLR with EfficientNet

Utility functions for image augmentations

Automatic data augmentation wrapper

SimCLR neural network for embeddings

Finding good initial-learning rate using LRFinder algorithm

Image classifier on top of SimCLR embeddings

Save the weights of neural network with embeddings

Classifier module

Classifier training code

Evaluation

Summary

Additional links and resources

Custom classifier on top of BERT-like Language Model - guide

TL;DR

Libraries and preparation

Understanding pre-trained language model input / outputs

Build custom classification head on top of LM

Preparing fast tokenizers to use with text dataset

Getting data for classification

Preparing reproducible training code with PyTorch Lightning

Hints and comments

Finding good starting learning rate for the model

Training the sentiment analysis classifier on PolEmo2.0

Using Trainer from PyTorch Lightning

Evaluating the trained model

Additional links and resources

Training RoBERTa from scratch - the missing guide

TL;DR

Getting the dataset

Preparing Polish Wikipedia dump

2. Stop using `.to(device)` that much: `zeros_like` / `ones_like` etc.

3. Register Buffer (a.k.a `nn.Module.register_buffer`)

4. Built-in `Identity()`

5. Pairwise distances: `torch.cdist`

6. Cosine similarity: `F.cosine_similarity`

7. Normalizing vectors: `F.normalize`

8. Linear layer + chunking trick (`torch.chunk`)

9. Masked select (`torch.masked_select`)

10. Conditional tensors with `torch.where`

11. Filling tensors with values at given positions (`Tensor.scatter`)

12. Image interpolation within network (`F.interpolate`)

13. Making grid of images (`torchvision.utils.make_grid`)

Adjusting the `run_language_modeling.py` script

Fast tokenizers support in `run_language_modeling.py`

This post is deprecated, see it's updated version:
https://zablo.net/blog/post/pytorch-resnet-mnist-jupyter-notebook-2021/