computer-vision on Marcin Zabłocki blog

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

Stroke Width Transform (SWT) for Python 3 & docker

Marcin Zabłocki — Mon, 11 Jan 2021 22:00:00 +0000

Back in august of 2017 I needed to use SWT algorithm to detect images that contain text in any form. At that time, there was no efficient enough implementation for my case, so I’ve decided to create Python bindings for libccv’s implementation of SWT. That was over 3 years ago - the solution got deprecated over time, now it’s officialy Python 3 era, with Docker all over the place. As the previous post still recieves some traffic and my GitHub repo is active, I have decided to update my solution to support newest stable libccv, Python 3.9 and Docker.

TL;DR

Stroke Width Transform for Python 3 - how to use a fast SWT implementation from libccv in Python 3.9 and/or Docker.

Prerequisites

Option A - local build:

g++ compiler
Python 3.9 (it will probably work with 3.6+ too)
cmake >= 3.17
libjpeg-dev, libpng-dev, libatlas-base-dev, libblas-dev installed
basic knowledge of C++
pybind11

Option B - docker image:

docker

SWT bindings for Python 3.9

Here is the how:

Download and compile libccv in the form of shared library
Implement bindings for the libccv using pybind11
Compile the library and use it

I just want to note, that I’m no C++ developer and I haven’t used pybind11 before, so in my opinion there is no experience with this toolset required.

Compling libccv for use in C++

There is actually not much to be done. I’ve cloned the original libccv repository and changed one flag in the makefile to make it work.

bash

CFLAGS := -O3 -fPIC -ffast-math -Wall $(CFLAGS)

See https://github.com/marrrcin/ccv/commit/5d945a6a02283880b4be4370d4dd6dd229402820

I recommend to just clone my fork from https://github.com/marrrcin/ccv/tree/stable :

bash

git clone --single-branch --branch stable --depth 1 https://github.com/marrrcin/ccv.git /ccv

Before the compilation starts, the following packages need to be installed in the system (it should work both on macOS and Linux):

bash

libjpeg-dev
libpng-dev
libatlas-base-dev
libblas-dev

After the required libraries are installed, I’ve just run the following command from the ./lib directory:

bash

./configure && make libccv.a

The compliation is rather a lengthy process if all of the optimization flags are enabled (which they are, by default). Also, when I run this compliation on macOS I experienced peak memory usage of 7GB, so keep this in mind. If you want to trade off library speed for the decreased compilation time/resources usage, replace the -O3 to either -O2, -O1 or -O0 do decrease / disable the compiler’s optimizations.

Implementing pybind11 bindings for SWT in libccv

I’ve created new cmake-based C++ project for a shared library and I’ve copied libccv.a and ccv.h files from the ./lib directory (of cloned libccv repo) into new folder within my project - also called lib.

The whole bindings for SWT consists of barely 41 lines of C++ code which I explain below.

csharp

// Include libccv headers to use stroke width transform

extern "C" {
    #include "lib/ccv.h"
}

// Include pybind11 headers to create Python bindings
#include <pybind11/pybind11.h>

namespace py = pybind11;

// Run SWT on the in-memory image (e.g. PNG or JPEG) and return the detected rectangles
py::list swt(char *bytes, int array_length){
    py::list result; // Initialize plain-Python List
    ccv_enable_default_cache();
    ccv_dense_matrix_t* image = 0;
    ccv_enable_default_cache();
    
    // Read in-memory image
    ccv_read((void*)bytes, &image, CCV_IO_GRAY | CCV_IO_ANY_STREAM, array_length);
    if(image != 0) {
        ccv_array_t *words = ccv_swt_detect_words(image, ccv_swt_default_params);
        if (words) {
            int i;
            for (i = 0; i < words->rnum; i++) {
                ccv_rect_t *rect = (ccv_rect_t *) ccv_array_get(words, i);
                py::dict item; // Initialize plain-Python dictionary and fill it with values
                item["x"] = rect->x;
                item["y"] = rect->y;
                item["width"] = rect->width;
                item["height"] = rect->height;
                result.append(item);
            }
            ccv_array_free(words);
        }
        ccv_matrix_free(image);
    }

    return result;
}

// Define pybind11 module
PYBIND11_MODULE(swt_python3, m) {
    m.doc() = "Python 3 compatible SWT (Stroke-Width-Transform) binding to libccv by Marcin Zabłocki";
    m.def("swt", &swt, "Apply Stroke-Width-Transform on input byte stream");
}

One key difference when compared to my old implementation is that with pybind11 I could use well-known Python structures such as lists and dictionaries directly from C++ code, making the whole “from C++ to Python” data passing a piece of cake 🎉.

In order to compile the bindings into a Python-usable package there are a few additional steps required.

Project needs to have CMakeLists.txt file in order for cmake to know how to compile it. In this file all of the necessary linker and pybind11 configurations are set:

python

cmake_minimum_required(VERSION 3.17)
project(swt_python3)

set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_FLAGS "-O3")
set(CMAKE_CXX_FLAGS_DEBUG "-g -O3")

find_library(CCV ccv lib)
find_package(JPEG REQUIRED)
find_package(PNG REQUIRED)
#find_package(PythonLibs) # Either uncomment this or specify `PYTHON_INCLUDE_DIRS` variable during build time
include_directories(${PYTHON_INCLUDE_DIRS})

add_subdirectory(pybind11)
pybind11_add_module(swt_python3)
target_sources(swt_python3 PRIVATE swt_python3.cpp)
target_link_libraries(swt_python3 PRIVATE ${CCV} ${JPEG_LIBRARIES} ${PNG_LIBRARIES})

Project needs to know where to find pybind11. It can be specified as another CMakeLists.txt file in the pybind11 folder within the project:

python

include(FetchContent)

FetchContent_Declare(
        pybind11
        URL   "https://github.com/pybind/pybind11/archive/v2.6.1.tar.gz"
)

FetchContent_MakeAvailable(pybind11)

That’s it. I compile the wrapper using the following command:

bash

cmake -DCMAKE_BUILD_TYPE=Release -DPYTHON_INCLUDE_DIRS=/usr/local/include/python3.9 . && make

It creates a file called swt_python3.cpython-39m-darwin.so (on macOS) or swt_python3.cpython-39-x86_64-linux-gnu.so (on Linux) which is Python-ready module to import!

TIP: in order to find PYTHON_INCLUDE_DIRS just run:

bash

python -c "from sysconfig import get_paths as gp; print(gp()['include'])"

Using stroke width transform in Python 3

Compiled binding is ready to be imported directly from Python code, so if you want to use it anywhere, just copy .so file into your project’s directory and call:

python

from swt_python3 import swt

Invoking SWT on any image is as simple as:

python

buffer = open("input.jpg", "rb").read()
swt_result: List[dict] = swt(buffer, len(buffer))
for item in swt_result:
    x, y, width, height = [item[key] for key in ("x", "y", "width", "height")]
    print(x, y, width, height)

Thanks to the pybind11 magic, the interface is much more plesant than last time.

Input

Output

SWT for Python 3 in Docker

As many modern deployments are Docker-based, I’ve decided to also include pre-built docker image which includes compiled and ready to use SWT wrapper with Python 3.9. The image includes an example.py script with the code that runs SWT on an image and saves the result into another one.

bash

docker pull marrrcin/swt-python3-ccv:20210111
docker run --rm -it --entrypoint bash marrrcin/swt-python3-ccv:20210111

Dockerfile: https://github.com/marrrcin/swt-python/blob/86a7553d1861f4dd7114c01b6ae512090d7ee14f/Dockerfile

Summary

I hope you liked it, please do not hesitate to create an issue or PR on GitHub if you find any issues.

Additional links & resources

full project on my github: https://github.com/marrrcin/swt-python
Docker Hub with the image: https://hub.docker.com/r/marrrcin/swt-python3-ccv
original paper about SWT: https://www.microsoft.com/en-us/research/publication/stroke-width-transform/

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

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

Stroke Width Transform algorithm for Python

Marcin Zabłocki — Fri, 11 Aug 2017 22:00:00 +0000

Stroke Width Transform (SWT) is a computer vision algorithm (actually, it’s an image operator) which can be used in the task of detecting text in images. This is a non-trivial task, especially for camera pictures, but SWT performs pretty well in this field.

This post is deprecated and the solution works only with Python 2.7 and old libccv versions.
For updated version (with Docker support) see: http://zablo.net/blog/post/stroke-width-transform-swt-python-3.9

TL;DR

How to wrap libccv computer vision library for using SWT in Python. Wrapper was built using SWIG.

Intro

My usecase for the SWT was the following: split big dataset of unknown images into ones that contain text (most likely documents / scans) and others. Easy, right? Yes, as long as you have fast implementation of this algorithm. Unfortunately, I needed to use Python, because the image filtering was only a part of bigger solution. I search here and there and of course, there are some implementations in Python, but they are not ready for production. One of them i.e took 5 minutes to process 2000px x 3000px image and it allocated gigabytes of memory. Just a single image. NOPE for that.

So… I searched further, but finnaly I haven’t found anything either working or performat. I decided to find any other language implementation of SWT and I’ve found libccv. Quick look at the repository - it’s written in C. OK, not my kind of language, but I gave it a try and surely, it will be fast. I’ve compiled example aplication which took image as an input and producted output in the form of 4-element tuples (x, y, width, height), each of which was an rectangle coordinates in which the text on the image is. It produced satysfying results, so I decided to use it from Python. How? Let’s dig into it!

Prerequisites

Python 2.7.* on Linux (tested on Ubuntu 16.04 x64)
Python packages: scikit-image, numpy, matplotlib (only for preview)
gcc compiler
SWIG
basic knowledge of C

SWT example breakdown from libccv

For starters, take a look at the code from https://github.com/liuliu/ccv/blob/07fc691c5344940751011c3af96d0ab202b1b4e6/bin/swtdetect.c The code comes from libccv repository and it’s an example of using SWT text detection. It is a console application, which takes file name as an argument, runs the algorithm and outputs coordinates of rectangles with text to standard output.

The probles with it are the following:

it is a console application
it takes file name as an input
it outputs to stdout

My requirements for the function (which I wanted to call from Python) were the following:

it is a standalone library / module
it takes blob of an image as an input (do not write anything to disk)
it outputs returning array-like object

SWT function wrapper in C

My requirements are pretty simple to transform to function header:

csharp

int* swt(char *bytes, int array_length, int width, int height);

Compiling the code above to plain english: take an array of bytes, it’s length (welcome to C world!), width and height of an image, do some magic and output array of numbers.

In order to implement and compile the wrapper without issues, just checkout libccv from it’s repository and put your files in libccv/lib folder. This approach is suggested by libccv author.

Implementation of SWT for Python wrapper:

csharp

#include "ccv.h"
#include <jpeglib.h>
#include "io/_ccv_io_libjpeg.c"
#include <sys/time.h>
#include <ctype.h>
#define SUCCESS 1
#define FAILURE 0

int* swt(char *bytes, int array_length, int width, int height){
    int *result_array;
    int status = FAILURE;
    ccv_enable_default_cache();
	ccv_dense_matrix_t* image = 0;

	FILE *stream;
	stream = fmemopen(bytes, array_length, "r");
	if(stream != NULL){
        int type = CCV_IO_JPEG_FILE | CCV_IO_GRAY;
        int ctype = (type & 0xF00) ? CCV_8U | ((type & 0xF00) >> 8) : 0;
        _ccv_read_jpeg_fd(stream, &image, ctype);

        if (image != 0)
        {
            ccv_array_t* words = ccv_swt_detect_words(image, ccv_swt_default_params);
            if (words)
            {
                int i;
                int result_idx = 1;
                result_array = (int*)malloc((4 * words->rnum + 1) * sizeof(int));
                result_array[0] = 4 * words->rnum;
                for (i = 0; i < words->rnum; i++)
                {
                    ccv_rect_t* rect = (ccv_rect_t*)ccv_array_get(words, i);
                    result_array[result_idx++] = rect->x;
                    result_array[result_idx++] = rect->y;
                    result_array[result_idx++] = rect->width;
                    result_array[result_idx++] = rect->height;
                }
                ccv_array_free(words);
                status = SUCCESS;
            }
            ccv_matrix_free(image);
        }
        ccv_drain_cache();
	}


    if(status != SUCCESS){
        result_array = (int*)malloc(1 * sizeof(int));
        result_array[0] = 0;
    }
    return result_array;
}

Explanation

There are some hacky things in there, I will explain them one by one.

csharp

FILE *stream;
stream = fmemopen(bytes, array_length, "r");

My function takes array of bytes as an input, but none of the ccv_read function overloads accepts this kind of data. I’ve found great fmemopen standard C function, which transforms char* buffer into file handler. But… none of the ccv_read function overloads accept this kind of data (AGAIN!). Yeah, but they read the JPEGs from files somehow!

That lead me to next part:

csharp

#include <jpeglib.h>
#include "io/_ccv_io_libjpeg.c"

In the _ccv_io_libjpeg.c file, there’s a function which accepts FILE* handler and reads JPEG image into libccv image format. I wanted to use it.

csharp

int type = CCV_IO_JPEG_FILE | CCV_IO_GRAY;
int ctype = (type & 0xF00) ? CCV_8U | ((type & 0xF00) >> 8) : 0;
_ccv_read_jpeg_fd(stream, &image, ctype);

Libccv needs to have grayscale image as an input for SWT, thus some bitwise magic to set proper flags.

csharp

int *results_array;
result_array = (int*)malloc((4 * words->rnum + 1) * sizeof(int));
result_array[0] = 4 * words->rnum;

Dynamic arrays were (and still are) such pain in C. SWT outputs 4 numbers for every match, so the size of the array should be 4 times bigger. The +1 is my hack to output everyhing in a single array and read it in Python. I used this one additional index to output the length of the whole array (I will drop this value in the upcoming SWIG wrapper, wait for it!).

csharp

if(status != SUCCESS){
    result_array = (int*)malloc(1 * sizeof(int));
    result_array[0] = 0;
}

When something goes wrong, I wanted to output something anyway, so I just create 1 element array, with 0, meaning, there is nothing to push to Python from the wrapper.

Python wrapper for SWT using SWIG

SWIG (http://www.swig.org/) seems to be really old, but it does the job pretty well (with really small amount of code). Plus it’s still active in 2017. It allows to wrap any C code into Python module really fast - that’s exactly what I wanted. You can install SWIG by using the following command:

powershell

sudo apt-get install swig

In order to build the wrapper, you need to create SWIG interface file (I named it ccvwrapper.i):

csharp

%module ccvwrapper

%{
#include "ccvwrapper.h"
%}

%typemap(out) int* swt {
	int i;
	$result = PyList_New($1[0]);
	for(i = 0; i < $1[0]; i++){
		PyObject *o = PyInt_FromLong($1[i+1]);
		PyList_SetItem($result, i, o);
	}
	free($1);
}

%include "ccvwrapper.h"

The most important path in here is the typemap delaration. This mapping converts my dynamic C array into Python list. When you want to return array from C code and use it in Python, you need to write this kind of code (this does not apply for basic types - they are automatically mapped to Python types). Header file ccvwrapper.h contains only my swt function declaration.

Building the wrapper

Just a simple bash script, following the convention from SWIG and libccv documentation. NOTE: remember to prefix output C library name with underscore _ ! SWIG requires it for building Python modules.

powershell

#!/bin/bash
swig -python ccvwrapper.i
gcc -fpic -c ccvwrapper.c ccvwrapper_wrap.c ccv_algebra.c ccv_basic.c ccv_cache.c ccv_classic.c ccv_io.c ccv_memory.c ccv_output.c ccv_resample.c ccv_sift.c ccv_swt.c ccv_transform.c ccv_util.c ./3rdparty/sha1/sha1.c -I/usr/include/python2.7

ld -shared ccvwrapper.o ccvwrapper_wrap.o ccv_algebra.o ccv_basic.o ccv_cache.o ccv_classic.o ccv_io.o ccv_memory.o ccv_output.o ccv_resample.o ccv_sift.o ccv_swt.o ccv_transform.o ccv_util.o sha1.o -ljpeg -o _ccvwrapper.so

cp ccvwrapper* ~/PycharmProjects/path_to_your_project
cp _ccvwrapper* ~/PycharmProjects/path_to_your_project

Build script uses swig command to compile interface file into two wrapepr files:

ccvwrapper_wrap.c
ccvwrapper.py

C file needs to be compiled along with the custom code and it’s dependencies. After compilation, we need to link everything into shared library with name _ccvwrapper.so. For the convenience I added two cp commands just to copy my wrapper into my project directory.

Using SWT Python wrapper

The wrapper is ready to be used. After putting all of the outputs from the build process into the project folder, you will be able to use the wrapper without any problems. Example code:

python

from __future__ import print_function
import ccvwrapper
import numpy as np
from skimage import draw
from skimage.io import imread, imshow, imsave
from matplotlib import pyplot as plt


def rectangle_perimeter(r0, c0, width, height, shape=None, clip=False):
    rr, cc = [r0, r0 + width, r0 + width, r0], [c0, c0, c0 + height, c0 + height]

    return draw.polygon_perimeter(rr, cc, shape=shape, clip=clip)


if __name__ == "__main__":
    image_name = "test_input.jpg"
    bytes = open(image_name, "rb").read()
    swt_result_raw = ccvwrapper.swt(bytes, len(bytes), 1024, 1360)
    swt_result = np.reshape(swt_result_raw, (len(swt_result_raw) / 4, 4))

    image = imread(image_name, as_grey=False)
    for x, y, width, height in swt_result:
        for i in xrange(0, 3): # just to make lines thicker
            rr, cc = rectangle_perimeter(y + i, x + i, height, width, shape=image.shape, clip=True)
            image[rr, cc] = (255, 0, 0)

    imshow(image)
    imsave("result.jpg", image)
    plt.show()

Results

Input

Output

You can download uncompressed input and output images here:

input
result

Summary

Seems pretty easy, right? With a little amount of work with, I was able to leverage power of SWT algorithm from Python with the speed of low-level C code.

Please comment if you encounter any problem and share if you like it!

Additional links & resources

full project on my github: https://github.com/marrrcin/swt-python
original paper about SWT: https://www.microsoft.com/en-us/research/publication/stroke-width-transform/
dataset of photos with text: http://www.liv.ic.unicamp.br/~minetto/datasets/text/IMAGES/CVPR/ORIG/

computer-vision on Marcin Zabłocki blog

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

Stroke Width Transform (SWT) for Python 3 & docker

TL;DR

Prerequisites

SWT bindings for Python 3.9

Compling libccv for use in C++

Implementing pybind11 bindings for SWT in libccv

Using stroke width transform in Python 3

SWT for Python 3 in Docker

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

Using ResNet for MNIST in PyTorch

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

TL;DR

Do not reinvent the wheel

Digging into the ResNet

Adjusting ResNet for MNIST

PyTorch training loop boilerplate

Imports:

Training loop for any PyTorch model (with evaluation)

Loading MNIST dataset and training the ResNet

Dataset

Training & evaluation

Summary

Additional links & resources

Stroke Width Transform algorithm for Python

This post is deprecated and the solution works only with Python 2.7 and old libccv versions.For updated version (with Docker support) see: http://zablo.net/blog/post/stroke-width-transform-swt-python-3.9

TL;DR

Intro

Prerequisites

SWT example breakdown from libccv

SWT function wrapper in C

Implementation of SWT for Python wrapper:

Explanation

Python wrapper for SWT using SWIG

Building the wrapper

Using SWT Python wrapper

Results

Summary

Additional links & resources

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

This post is deprecated and the solution works only with Python 2.7 and old libccv versions.
For updated version (with Docker support) see: http://zablo.net/blog/post/stroke-width-transform-swt-python-3.9