Extracting rich embedding features from COCO pictures using PyTorch and ResNeXt-WSL
By: Gianmario Spacagna
Tags: Resnext, K Neighbors, Pytorch, Embeddings, Computer Vision

Extracting rich embedding features from COCO pictures using PyTorch and ResNeXt-WSL

Photo by Cosmic Timetraveler on Unsplash

In this tutorial, I will show you how to leverage a powerful pre-trained convolution neural network to extract embedding vectors that can accurately describe any kind of picture in an abstract latent feature space.

I will show some examples of using ResNext-WSL on the COCO dataset using the library PyTorch and other conventional tools from the PyData stack.

Why ResNext-WSL?

ResNeXt is the evolution of the well famous ResNet model that adds an additional dimension on top of it called the “cardinality” dimension. Through this improvement, the authors managed to beat the benchmark of ILSVRC classification task with a 15% improvement. Although better models were developed in the last couple of years, it still stands on the top leaderboard of the ImageNet task, with 85.4% top-1 accuracy.

ResNext block from paperswithcode

The weakly supervised learning (WSL) version was developed by the Facebook AI Research (FAIR) group, but it was trained on 3.5 billion public Instagram pictures in order to predict around 8000 hashtags sourced by the users.

This novel approach made ResNeXt-WSL my favorite choice among many other publicly available. I believe that a model trained on a broader variety of pictures can encapsulate many different topics of real-world scenarios while other pre-trained models were mostly optimized solely for their benchmark task.

Below are some facts about the training of this massive model (source):

synchronous stochastic gradient descend on 336 GPUs across 42 machines with mini-batches of 8064 images

153 billion multiply-add FLOPs and 829 million model parameters for the largest capacity model

22 days long training for the smaller capacity model (32x16d)

The model was trained with different capacities in order to reduce its size and inference time.

Different ResNext models and capacities from PyTorch Hub

After training the model on the Instagram hashtag classification task, the last layer was replaced with the ImageNet classes and fine-tuned on that task. Clearly training those models is not something every organization can afford.

Luckily for us, a simplified version was released in the PyTorch hub, trained on “only” 940 million public images and 1.5k hashtags. For this tutorial, we choose the smallest version of ResNext-WSL (32x8d) since it is still accurate enough on the public benchmarks but it is less memory and computationally expensive.

Why COCO dataset?

Common Objects in COntext ( COCO) is a large-scale object detection, segmentation, and captioning dataset, widely used as a benchmark for many machine learning tasks.

COCO dataset logo from COCO website

For this tutorial, we would focus on the Detection 2017 dataset (validation fold) consisting of 5000 annotated pictures. The dataset satisfies a few desired properties:

Contains pictures from disparate context and kind of objects

The pictures are realistic with different shapes and resolutions

Each picture may contain multiple annotations

The annotations are grouped into two levels hierarchical taxonomy

There are data loaders and libraries for COCO already implemented in Python and PyTorch

The 5000 pictures are annotated with 80 categories grouped in 12 supercategories:

Distribution of categories in COCO 2017 validation set

And below are a sample of random pictures with their annotations:

Examples of annotated pictures from COCO detection 2017

As we can see the pictures in the COCO dataset do represent different real-world scenarios, even though they are skewed with a lot of pictures containing people.

Feature Extraction

The ResNeXt traditional 32x4d architecture is composed by stacking multiple convolutional blocks each composed by multiple layers with 32 groups and a bottleneck width equal to 4. That is the first convolution layer with 64 filters is parallelized in 32 independent convolutions with only 4 filters each.

ResNet versus ResNext architecture comparison from PyTorch Hub

The smallest WSL model in the PyTorch hub uses a 32x8d architecture (bottleneck width = 8) instead. Regardless of those architecture details, what matters for us is the last average pool layer that would concatenate all of those filters into a single dimensional array of size 2048.

The concatenated features are then supposed to be fed to the output softmax layer predicting the 1000 classes of ImageNet. Since we are not interested in the class predictions, we will drop the softmax layer and use the array of the average pool as the embedding features for our pictures.

The embedding-only model will have the following size:

Embedding model size

The first thing to do in order to be usable is to pre-process the input pictures in the format the model would expect. The preprocessing consists of:

Scaling to 256×256

Centering crop to 224×224

Normaliing with mean = [0.485, 0.456, 0.406] and stdev = [0.229, 0.224, 0.225]

I have used a Google Colab environment with GPU runtime to perform the loading, preprocessing, and inference on the 5000 pictures in around 10 minutes.

We now have obtained a matrix of (5000, 2048) float numbers.

Feature exploration

Let’s dive into the feature space and extract some insights.

For example, if we consider a sample picture with a dog:

Sample picture from pytorch github repository

It would be embedded with the following values:

Distribution of embedding values for the sample picture

We can notice that by the design of the neural network we only get non-negative values without any upper bound.

Since not all of the 2048 features will contribute the same way, we can train a random forest classifier on the COCO multi-label annotation task in order to estimate the feature importance.

Let’s now zoom into the most important features extracted from the random forest classifier, whose first 25 are: [ 12, 3, 2011, 1151, 28, 1715, 1004, 583, 479, 1377, 1505, 1690, 1028, 1831, 1704, 466, 634, 1329, 2003, 176, 485, 1066, 1669, 445, 1227].

Correlation matrix between pairs of features

Correlation scatter plots between pairs of features

We can derive a couple of properties:

The feature space is slightly sparse (many feature values are close to 0) The feature values for each dimension follow an exponential distribution on our dataset

The measure the first point we can count what percentage of the dimensions are used by each picture (value > 0.01):

Distribution of the percentage of dimensions activated in each picture

Only 40% of dimensions are used in average by each picture, where we can assume the higher the value the more the topics embedded in the picture.

In order to confirm the exponential distribution hypothesis we can plot the mean of each dimension against its standard deviation across the dataset:

Mean vs standard deviation of the 2048 dimensions on COCO pictures

The proportionality between mean and standard deviation gives us a hint that the dimensions tend to follow an exponential-like distribution on our dataset.

3D visualization

We can apply a few dimensionality reduction techniques to reduce the embedding space into 3 dimensions and use the supercategory to highlight the pictures that belong to the same taxonomy.

We will focus on 3 unsupervised techniques: PCA, t-SNE, and UMAP. For a more exhaustive comparison of their differences, I recommend the following reading: Dimensionality Reduction for Data Visualization: PCA vs TSNE vs UMAP vs LDA.

3D visualization using PCA

3D visualization using TSNE

3D visualization using UMAP

To make the exploration more fun I have deployed a TensorBoard using a sample of pictures in a sprite image:

Sprite image for interactive visualization in TensorBoard

And the TSNE projection looked like this:

TSNE projection animation

The interactive projector of the TensorBoard is a powerful tool for exploring the embedding space and spot correlations and interesting manifolds of points.

Nearest Neighbors

What can we actually do now with this embedding space? A relevant application would be to find similar pictures. Thus, I have selected a few random pictures and plotted the 15 closest ones based on the cosine similarity in the 2048-dimensional embedding space.

The results are astonishing:

“Bikes” neighborhood

“Tennis players” neighborhood

“Living room with sofas” neighborhood

“People skiing” neighborhood

“Cats on a laptop” neighborhood

“Food on a plate” neighborhood

“Large animals exhibition” neighborhood

Not just the embedding space is able to accurately group pictures that do are very similar in terms of objects and context, but it can also correctly discriminate the entropy on each kind of picture based on the cosine distance. For instance, the neighborhoods of “tennis players” and “people skiing” have a much lower distance from the query picture compared to the other more entropic neighborhoods.

Conclusions

In this tutorial, we have seen how easily, and without having to train any model, we can leverage the state-of-the-art to extract rich embedding features from pictures of different natures. We have proven the ability of the embedding space to capture a lot of fine details including both the objects and the surrounding context. The most salient part was showing that the embedding features can accurately find similar pictures. The measure of similarity and the latent manifold structures can enable a lot of downstream applications.

In the next articles, we will see how to discover and define those latent topics represented by the manifolds in the embedding space, to learn how to cluster those pictures, and to learn advanced averaging techniques for document embedding.

Stay tuned!

You can find the code and the notebooks at https://github.com/gm-spacagna/docem.

If you are interested in extracting embedding features from text data, you can check out this other article: Embedding billions of text documents using Tensorflow Universal Sentence Encoder and Spark EMR.