Extracting Fixed-Size Feature Vectors

The "Why": Preparing for Fusion

Your project's goal is to combine signals from multiple data types (images, text, metadata). A "fusion model" can't directly process a JPEG image and a string of text. It needs these inputs to be in a consistent, numerical format. The process involves taking each unimodal model (like your image CNN) and using it to generate a feature vector, which is then fed into a fusion module.

This diagram illustrates the late fusion strategy we are building towards. Each input modality is processed by its own "Unimodal model" to produce an intermediate representation (a feature vector). These vectors are then combined in the "Fusion module" for a final prediction. Today, we build the first part of this pipeline: converting an image into its feature vector.

A feature vector, also called an embedding, is simply a list of numbers that acts as a dense, numerical summary of the input. For an image, this vector encodes information about textures, shapes, colors, and object arrangements that the model learned during training.

The "Where": Finding Features in a CNN

In Module 4, you fine-tuned a Convolutional Neural Network (CNN) to classify meal quality. A CNN has two main parts: a feature extraction backbone (convolutional and pooling layers) and a classification head (fully connected layers).

A typical CNN architecture. The initial layers are responsible for feature extraction, learning to identify progressively more complex patterns. The final fully connected layers perform the classification. The feature vector we want is the output of the feature extraction stage, right before it's passed to the classifier.

The final nn.Linear layer of your ResNet model was trained specifically to map features to your three quality classes ('Good', 'Adequate', 'Poor'). This layer is too task-specific. The layers before it, however, have learned a much richer, more general-purpose representation of the image. Our goal is to grab the output tensor from the layer just before this final classifier. For a ResNet model, this is typically the output of the average pooling layer (avgpool). This output is our feature vector.

To do this, we need a way to "intercept" the data as it flows through the network during a forward pass. There are two primary ways to achieve this in PyTorch.

Method 1: PyTorch Forward Hooks

The first method uses a powerful, low-level PyTorch feature called a forward hook. A hook is a function that you can "register" to a specific nn.Module (a layer). PyTorch will automatically execute your hook function right after that layer's forward method completes, giving you access to the layer's input and, crucially, its output.

This video from the mildlyoverfitted channel provides an excellent explanation of what hooks are and how they work under the hood.

The following article from STACC provides a complete, practical example of using a forward hook to extract a feature vector from a ResNet-18 model. We will adapt this logic for our fine-tuned ResNet-50.

Implementing Feature Extraction with a Hook

Let's put this into practice. The script below defines a class, FeatureExtractor, that encapsulates this logic. It loads your fine-tuned model and uses a forward hook to extract the feature vector.

Before you run this:

1. Make sure you have the fine-tuned ResNet model you saved in Module 4 (e.g., meal_quality_cnn.pt).
2. Remember that the image transforms (resizing, normalization) must be the exact same as the ones you used for validation/testing when you trained the model.

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image

class FeatureExtractor:
    def __init__(self, model_path: str, num_classes: int, device: str = 'cuda'):
        """
        Initializes the feature extractor.

        Args:
            model_path (str): Path to the fine-tuned model checkpoint.
            num_classes (int): Number of output classes the model was trained on.
            device (str): The device to run the model on ('cuda' or 'cpu').
        """
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        
        # Load a pretrained ResNet-50 structure
        self.model = models.resnet50()
        
        # Get the number of input features for the classifier
        num_ftrs = self.model.fc.in_features
        
        # Replace the final layer with one matching our number of classes
        self.model.fc = nn.Linear(num_ftrs, num_classes)
        
        # Load the saved weights
        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
        
        # Set the model to evaluation mode
        self.model.eval()
        
        # Move model to the specified device
        self.model.to(self.device)

        # Define the layer to hook into
        self.layer = self.model._modules.get('avgpool')
        
        # Define the image transformations
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def _hook_function(self, module, input, output):
        # Squeeze to remove the batch dimension, then copy to CPU
        self.embedding = output.squeeze().detach().cpu()

    def extract(self, image_path: str) -> torch.Tensor:
        """
        Extracts a feature vector from a single image.

        Args:
            image_path (str): Path to the input image.

        Returns:
            torch.Tensor: The extracted feature vector.
        """
        img = Image.open(image_path).convert('RGB')
        transformed_img = self.transform(img).unsqueeze(0).to(self.device)

        # Register the hook
        h = self.layer.register_forward_hook(self._hook_function)

        # Perform a forward pass
        with torch.no_grad():
            self.model(transformed_img)

        # Remove the hook
        h.remove()

        return self.embedding

# --- Main execution block ---
if __name__ == "__main__":
    # Path to your fine-tuned model from Module 4
    MODEL_CHECKPOINT = 'path/to/your/meal_quality_cnn.pt'
    # Path to an image you want to process
    IMAGE_FILE = 'path/to/your/meal_image.jpg'
    
    # Initialize the extractor
    extractor = FeatureExtractor(model_path=MODEL_CHECKPOINT, num_classes=3)
    
    # Extract the feature vector
    feature_vector = extractor.extract(IMAGE_FILE)
    
    print("--- Feature Vector ---")
    print(f"Shape: {feature_vector.shape}")
    print(f"Vector (first 10 elements): {feature_vector[:10]}...")
    
    # The output shape for a ResNet-50's avgpool layer is torch.Size([2048])

Code Walkthrough:

• __init__: Loads your fine-tuned model, sets it to eval() mode (crucial for consistent results), and defines the image transforms.
• _hook_function: This is the function that does the work. It takes the output from the avgpool layer, removes the batch dimension (.squeeze()), and copies it.
• extract: This method takes an image path, applies the transforms, registers the hook, runs the forward pass (inside a torch.no_grad() context for efficiency), and then removes the hook.

Method 2: The torchvision.models.feature_extraction Utility

While forward hooks are powerful and flexible, PyTorch provides a more modern, high-level utility that often makes this task simpler: create_feature_extractor. This function symbolically traces your model, rebuilds it to output the activations from specified layers, and wraps it in a new nn.Module.

The official Torchvision documentation explains this utility well.

Here is how you would use create_feature_extractor to achieve the same goal. Notice how much simpler the code is, as you don't need to manage hooks manually.

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from torchvision.models.feature_extraction import create_feature_extractor
from PIL import Image

# Path to your fine-tuned model and an image
MODEL_CHECKPOINT = 'path/to/your/meal_quality_cnn.pt'
IMAGE_FILE = 'path/to/your/meal_image.jpg'
NUM_CLASSES = 3

# 1. Load your original fine-tuned model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
base_model = models.resnet50()
base_model.fc = torch.nn.Linear(base_model.fc.in_features, NUM_CLASSES)
base_model.load_state_dict(torch.load(MODEL_CHECKPOINT, map_location=device))

# 2. Create the feature extractor model
# We specify the node 'avgpool' as the output we want.
feature_extractor = create_feature_extractor(base_model, return_nodes=['avgpool']).to(device)
feature_extractor.eval() # Set to evaluation mode

# 3. Prepare the image (same transforms as before)
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
img = Image.open(IMAGE_FILE).convert('RGB')
transformed_img = transform(img).unsqueeze(0).to(device)

# 4. Extract features
with torch.no_grad():
    features = feature_extractor(transformed_img)

# The output is a dictionary where keys are the node names
feature_vector = features['avgpool'].squeeze()

print("--- Feature Vector (using create_feature_extractor) ---")
print(f"Shape: {feature_vector.shape}")
print(f"Vector (first 10 elements): {feature_vector.cpu()[:10]}...")

This second method is generally recommended for its simplicity and clarity. The output is a dictionary, making it easy to extract from multiple layers if needed.

Conclusion

In this lesson, you've learned a fundamental skill for multimodal deep learning: converting raw data into a fixed-size feature vector. This process transforms a high-dimensional input like an image into a dense numerical representation that a downstream model can work with.

Your key takeaways are:

• Feature vectors are the common language for multimodal fusion, representing complex data in a simple numerical format.
• In a CNN, the best features for transfer tasks are typically extracted from layers just before the final classifier.
• You must always set your model to evaluation mode (model.eval()) before extracting features to ensure deterministic and correct output.
• PyTorch offers two main ways to do this: low-level forward hooks for maximum control, and the high-level create_feature_extractor utility for simplicity and convenience.

You now have the toolset to generate features from your entire image dataset. In our next lesson, we will tackle the other half of our fusion problem: text. You'll learn about the Transformer architecture and how to use a powerful pretrained model, BERT, to extract a feature vector from the parent complaint texts.