Feature Extraction

Convolutional networks used for image classification consist of two main components: a backbone composed of convolution and pooling layers, followed by a classification head made up of fully-connected layers. Broadly speaking, the convolutional backbone extracts generic features from the input image, while the classification head interprets these features to make a prediction. Transfer learning takes advantage of this architecture to repurpose the convolutional backbone of a pretrained network. Instead of training an entire model from scratch, the pre-trained backbone is reused to extract features from new data. Then, a new classifier is trained on top of these features. This is commonly referred to as feature extraction.

In this tutorial, we will explain how to use feature extraction on the cats-vs-dogs dataset. We will download a pretrained model from the TorchVision library, get the convolutional backbone from the model, and use it to extract features from images of cats and dogs. We will then train a new classifier on top of these features.

Analysis of MobileNetV3

Let’s take a look at a model that has been pretrained on the ImageNet dataset. There are plenty of models to choose from in the TorchVision library. In this tutorial, we use MobileNetV3 for its speed and efficiency.

from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights

weights = MobileNet_V3_Large_Weights.DEFAULT

pretrained = mobilenet_v3_large(weights=weights)

We inspect MobileNetV3 architecture using the TorchInfo library, which provides a detailed summary of the model’s layers, output shapes, number of parameters, and more.

Note

If you followed the instructions to set up your environment, you already have torchinfo installed.

To have a breakdown of the output shapes, we need to specify the expected input size of the model. MobileNetV3 expects images of size 224x224 with three color channels (RGB). The batch size is set to 1 for simplicity. We also set the depth parameter to limit the level of detail in the summary.

Show code cell source

Hide code cell source

from torchinfo import summary

summary(pretrained, input_size=(1, 3, 224, 224), depth=1, col_width=16, col_names=["input_size", "output_size"], row_settings=["var_names"])

============================================================================================
Layer (type (var_name))                                      Input Shape      Output Shape
============================================================================================
MobileNetV3 (MobileNetV3)                                    [1, 3, 224, 224] [1, 1000]
├─Sequential (features)                                      [1, 3, 224, 224] [1, 960, 7, 7]
├─AdaptiveAvgPool2d (avgpool)                                [1, 960, 7, 7]   [1, 960, 1, 1]
├─Sequential (classifier)                                    [1, 960]         [1, 1000]
============================================================================================
Total params: 5,483,032
Trainable params: 5,483,032
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 216.62
============================================================================================
Input size (MB): 0.60
Forward/backward pass size (MB): 70.46
Params size (MB): 21.93
Estimated Total Size (MB): 92.99
============================================================================================

As indicated above, MobileNetV3 consists of three modules.

• features - A series of convolutional and pooling layers that extract the input image into feature maps.

• avgpool - A layer that averages the channels of the feature maps (3D) to produce a feature vector (1D).

• classifier - A series of fully-connected layers that predicts the class logits from the feature vector.

Classification head

Let’s visualize the classification head of MobileNetV3. The summary shows that the first layer expects a vector of size 960, while the last layer produces a vector of size 1000. We take note of the input size because we will need this information to build a new classifier on top of the extracted features.

Show code cell source

Hide code cell source

summary(pretrained.classifier, input_size=(1, 960), col_names=["input_size", "output_size"])

==========================================================================================
Layer (type:depth-idx)                   Input Shape               Output Shape
==========================================================================================
Sequential                               [1, 960]                  [1, 1000]
├─Linear: 1-1                            [1, 960]                  [1, 1280]
├─Hardswish: 1-2                         [1, 1280]                 [1, 1280]
├─Dropout: 1-3                           [1, 1280]                 [1, 1280]
├─Linear: 1-4                            [1, 1280]                 [1, 1000]
==========================================================================================
Total params: 2,511,080
Trainable params: 2,511,080
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 2.51
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.02
Params size (MB): 10.04
Estimated Total Size (MB): 10.07
==========================================================================================

Convolutional backbone

Let’s visualize the convolutional backbone of MobileNetV3. We notice that the expected input has 3 channels, while the produced output has 960 channels. This aligns with the classification head.

Show code cell source

Hide code cell source

summary(pretrained.features, input_size=(1, 3, 224, 224), depth=1, col_names=["input_size", "output_size"])

===============================================================================================
Layer (type:depth-idx)                        Input Shape               Output Shape
===============================================================================================
Sequential                                    [1, 3, 224, 224]          [1, 960, 7, 7]
├─Conv2dNormActivation: 1-1                   [1, 3, 224, 224]          [1, 16, 112, 112]
├─InvertedResidual: 1-2                       [1, 16, 112, 112]         [1, 16, 112, 112]
├─InvertedResidual: 1-3                       [1, 16, 112, 112]         [1, 24, 56, 56]
├─InvertedResidual: 1-4                       [1, 24, 56, 56]           [1, 24, 56, 56]
├─InvertedResidual: 1-5                       [1, 24, 56, 56]           [1, 40, 28, 28]
├─InvertedResidual: 1-6                       [1, 40, 28, 28]           [1, 40, 28, 28]
├─InvertedResidual: 1-7                       [1, 40, 28, 28]           [1, 40, 28, 28]
├─InvertedResidual: 1-8                       [1, 40, 28, 28]           [1, 80, 14, 14]
├─InvertedResidual: 1-9                       [1, 80, 14, 14]           [1, 80, 14, 14]
├─InvertedResidual: 1-10                      [1, 80, 14, 14]           [1, 80, 14, 14]
├─InvertedResidual: 1-11                      [1, 80, 14, 14]           [1, 80, 14, 14]
├─InvertedResidual: 1-12                      [1, 80, 14, 14]           [1, 112, 14, 14]
├─InvertedResidual: 1-13                      [1, 112, 14, 14]          [1, 112, 14, 14]
├─InvertedResidual: 1-14                      [1, 112, 14, 14]          [1, 160, 7, 7]
├─InvertedResidual: 1-15                      [1, 160, 7, 7]            [1, 160, 7, 7]
├─InvertedResidual: 1-16                      [1, 160, 7, 7]            [1, 160, 7, 7]
├─Conv2dNormActivation: 1-17                  [1, 160, 7, 7]            [1, 960, 7, 7]
===============================================================================================
Total params: 2,971,952
Trainable params: 2,971,952
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 214.11
===============================================================================================
Input size (MB): 0.60
Forward/backward pass size (MB): 70.44
Params size (MB): 11.89
Estimated Total Size (MB): 82.93
===============================================================================================

To retrieve the exact size of the output, we can pass a random tensor of shape (…, 3, 224, 224) through the backbone. We then inspect the shape of the resulting tensor.

import torch

batch = torch.randn(1, 3, 224, 224)

model_device = next(pretrained.parameters()).device
batch = batch.to(model_device)

with torch.inference_mode():
    output = pretrained.features(batch)

print('Feature shape:', tuple(output.shape[1:]))

Feature shape: (960, 7, 7)

Standalone feature extractor

There are two common ways to use a frozen backbone.

• Standalone extraction. Every image is run through the backbone once, the resulting tensors are saved, and a classifier is trained from the saved features. This is fast because the expensive backbone is not executed every epoch, but it does not allow for random augmentation of the input images.

• End-to-end extraction. The frozen backbone is combined with a trainable head in a single model. Images pass through the backbone every epoch, which costs more compute but permits new random augmentations on each pass.

We will cover the standalone approach in this tutorial.

Show code cell source

Hide code cell source

import torch
import torch.nn.functional as F
from torch import nn, optim
from torch.utils import data
from torchvision.datasets import ImageFolder
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

Dataset and splits

We start by loading the cats-vs-dogs dataset with the preprocessing attached to MobileNetV3. We create three stratified subsets for training, validation, and testing.

Show code cell source

Hide code cell source

data_path = '.data/cats_vs_dogs/PetImages'

weights = MobileNet_V3_Large_Weights.DEFAULT
preprocess = weights.transforms()

print("Weight recipe:", weights)

# Full dataset with MobileNetV3 preprocessing
dataset = ImageFolder(
    data_path,
    transform=preprocess,
    target_transform=lambda x: torch.tensor(x, dtype=torch.float),
)

# Training split
all_idx = list(range(len(dataset)))
train_idx, holdout_idx = train_test_split(
    all_idx,
    stratify=dataset.targets,
    test_size=0.30,
    random_state=42,
)

# Validation and test splits
holdout_targets = [dataset.targets[i] for i in holdout_idx]
valid_idx, test_idx = train_test_split(
    holdout_idx,
    stratify=holdout_targets,
    test_size=0.50,
    random_state=42,
)

train_ds = data.Subset(dataset, train_idx)
valid_ds = data.Subset(dataset, valid_idx)
test_ds = data.Subset(dataset, test_idx)

print(f"Train: {len(train_ds)}, validation: {len(valid_ds)}, test: {len(test_ds)}")

Weight recipe: MobileNet_V3_Large_Weights.IMAGENET1K_V2
Train: 17497, validation: 3750, test: 3750

See also

The Dataset from files tutorial explains how to download and validate the cats-vs-dogs dataset.

Extracting features

Let’s extract features from the cats-vs-dogs dataset by running the convolutional backbone over the whole dataset. We define a function extract_features that takes a model and a dataset as input, and returns a TensorDataset with the extracted features and the corresponding labels.

Show code cell source

Hide code cell source

from tqdm import tqdm

def extract_features(model: nn.Module,
                     dataset: data.Dataset,
                     batch_size: int = 64) -> data.TensorDataset:
    
    previous_device = next(model.parameters()).device
    was_training = model.training

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

    loader = data.DataLoader(dataset, batch_size=batch_size, shuffle=False)
    features, labels = [], []

    with torch.inference_mode():
        for images, targets in tqdm(loader, desc="Extracting features"):
            images = images.to(device)
            feats = model(images)
            features.append(feats.cpu())
            labels.append(targets)

    model.to(previous_device).train(was_training)
    
    return data.TensorDataset(torch.cat(features), torch.cat(labels))

Feature extraction runs the convolutional backbone over every image and can take several minutes on a CPU. The three resulting datasets could be saved with torch.save and reused.

pretrained = mobilenet_v3_large(weights=weights)

extracted_train_ds = extract_features(pretrained.features, train_ds)
extracted_valid_ds = extract_features(pretrained.features, valid_ds)
extracted_test_ds = extract_features(pretrained.features, test_ds)

Extracting features: 100%|██████████| 274/274 [08:08<00:00,  1.78s/it]
Extracting features: 100%|██████████| 59/59 [01:39<00:00,  1.69s/it]
Extracting features: 100%|██████████| 59/59 [01:43<00:00,  1.76s/it]

Let’s print the size of the extracted features. MobileNetV3 backbone produces spatial feature maps with shape (960, 7, 7) for each 224×224 input. We keep this representation for the classifier that is going to be trained after feature extraction. MobileNetV3 normally applies adaptive average pooling on the feature maps, but we do not use this pooling here.

Show code cell source

Hide code cell source

print('Extracted features (train):     ', tuple(extracted_train_ds.tensors[0].shape))
print('Extracted features (validation):', tuple(extracted_valid_ds.tensors[0].shape))
print('Extracted features (test):      ', tuple(extracted_test_ds.tensors[0].shape))

Extracted features (train):      (17497, 960, 7, 7)
Extracted features (validation): (3750, 960, 7, 7)
Extracted features (test):       (3750, 960, 7, 7)

Defining a new classifier

Now that we have the extracted features, we can train a new classifier on top of them. We will use a simple feedforward neural network with two hidden layers. The input size of the network should match the size of the extracted features, while the output size is a single value representing the logit for the positive class.

class Classifier(nn.Module):
    
    def __init__(self):
        super().__init__()
        input_dim = 960 * 7 * 7
        self.fc = nn.Linear(input_dim, 256)
        self.dropout = nn.Dropout(0.5)
        self.out = nn.Linear(256, 1)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.fc(x)
        x = torch.relu(x)
        x = self.dropout(x)
        x = self.out(x)
        return torch.squeeze(x, -1)

Note

The model output is squeezed so that it has shape (batch_size,) instead of (batch_size, 1). At the same time, labels are converted to float tensors during preprocessing. Together, these steps prevent shape and type mismatch errors when using BCEWithLogitsLoss as a loss function.

Training the classifier

Let’s import Trainer from the training.py file and BinaryAccuracy from TorchEval.

from training import Trainer
from torcheval.metrics import BinaryAccuracy

Training the head is much faster than training the complete CNN because it reads cached tensors rather than recomputing backbone features. The cache has shifted the cost to a one-time extraction step and additional storage.

train_loader = data.DataLoader(extracted_train_ds, batch_size=64, shuffle=True)
valid_loader = data.DataLoader(extracted_valid_ds, batch_size=128, shuffle=False)
test_loader = data.DataLoader(extracted_test_ds, batch_size=128, shuffle=False)

torch.manual_seed(42)

model = Classifier()
loss_fn = nn.BCEWithLogitsLoss()

optimizer = optim.Adam(model.parameters(), lr=0.001, amsgrad=True)
epochs = 5

trainer = Trainer()
trainer.set_metrics(accuracy=BinaryAccuracy())

history = trainer.fit(model, train_loader, loss_fn, optimizer, epochs, valid_loader)

===== Training on cpu device =====

Epoch  1/5: 100%|██████████| 274/274 [00:31<00:00,  8.67it/s, accuracy=0.9880, train_loss=0.0869, valid_loss=0.0490]
Epoch  2/5: 100%|██████████| 274/274 [00:23<00:00, 11.84it/s, accuracy=0.9864, train_loss=0.0227, valid_loss=0.0423]
Epoch  3/5: 100%|██████████| 274/274 [00:23<00:00, 11.60it/s, accuracy=0.9877, train_loss=0.0067, valid_loss=0.0550]
Epoch  4/5: 100%|██████████| 274/274 [00:22<00:00, 12.45it/s, accuracy=0.9904, train_loss=0.0072, valid_loss=0.0565]
Epoch  5/5: 100%|██████████| 274/274 [00:22<00:00, 12.25it/s, accuracy=0.9880, train_loss=0.0062, valid_loss=0.0786]

Let’s plot the training loss and monitor the validation loss and accuracy. The validation curves can help us make informed decisions about model selection and early stopping; the test set is not involved in this process.

Show code cell source

Hide code cell source

plt.figure(figsize=(10, 5), tight_layout=True)
plt.subplot(1, 2, 1)
plt.plot(history['train_loss'], label='Train loss')
plt.plot(history['valid_loss'], label='Validation loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.subplot(1, 2, 2)
plt.plot(history['accuracy'], label='Validation accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

You should notice that the training loss continues to improve while the validation loss deteriorates. This is a sign that the head is overfitting the cached training features, despite the dropout layer added to the classifier. Data augmentation may help mitigate overfitting here, but it would require a switch to the end-to-end approach, which is not covered in this tutorial.

Final evaluation

After all choices are fixed using the training and validation data, we can evaluate the classifier once on the cached test features.

Show code cell source

Hide code cell source

ans = trainer.eval(model, test_loader)

print(f"Test accuracy: {ans['accuracy']:.2%}")

Test accuracy: 98.88%

Let’s also inspect a few examples from the test set. We use the test indices explicitly so the images correspond to the held-out subset, and move both the backbone and features to the appropriate devices.

Show code cell source

Hide code cell source

device = next(model.parameters()).device
pretrained.features.to(device).eval()

torch.manual_seed(14)
n = 9
positions = torch.randint(0, len(test_idx), (n,))
raw_indices = [test_idx[position] for position in positions.tolist()]

raw_data = ImageFolder(data_path)
images, labels = zip(*[raw_data[i] for i in raw_indices])
batch = torch.stack([preprocess(img) for img in images])

with torch.inference_mode():
    batch = batch.to(device)
    features = pretrained.features(batch)
    outputs = model(features)

plt.figure(figsize=(9, 9), tight_layout=True)
for i in range(n):
    plt.subplot(3, 3, i + 1)
    plt.imshow(images[i])
    plt.title(f'Prediction: {"Dog" if outputs[i] > 0 else "Cat"}')
    plt.axis('off')
plt.show()

Summary

In this tutorial, we learned how to use feature extraction to repurpose a pretrained convolutional backbone for a new classification task. We downloaded a pretrained MobileNetV3 from the TorchVision library, extracted features from the cats-vs-dogs dataset, and trained a new classifier on top of these features. The classifier achieved an impressive accuracy of 99% on the test set, which is not surprising given the simplicity of the dataset and its similarity to the ImageNet dataset.

FAQ

• Why did not we reuse the pretrained classification head?

  • Because the ImageNet classification head is trained to predict 1,000 classes.

• Should we always reuse the entire convolutional backbone?

  • Earlier layers often represent broadly useful visual patterns, while deeper features are more specialized. Whether the full backbone transfers well depends on the distance between source and target domains. If the new dataset differs significantly from the dataset that the original model was trained on, it may be better to use only the first few layers of the model for feature extraction, rather than using the entire convolutional backbone.