PyTorch - Feature Extraction in Convents

PyTorch - Feature Extraction with Convolutional Neural Networks (ConvNets)

Convolutional Neural Networks (ConvNets or CNNs) are powerful for tasks like image recognition because they automatically learn hierarchical features from data. Feature extraction is the process of using parts of a pre-trained or custom ConvNet to obtain these learned features, which can then be used for other tasks like transfer learning, image similarity, or as input to other machine learning models.

We’ll explore how to extract features using PyTorch, primarily focusing on leveraging pre-trained models from torchvision.models.

Step 1: Import Necessary Libraries

We need torch, torch.nn for neural network modules, and torchvision.models for pre-trained architectures, and torchvision.transforms for image preprocessing.

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image # For loading an example image

Set up the device for computation:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

Step 2: Load a Pre-trained Model

torchvision.models provides many pre-trained models. We’ll use ResNet-18 as an example. The weights parameter allows us to load pre-trained weights. Using SomeModel_Weights.DEFAULT is the modern way to get the best available pre-trained weights.

# Load a pre-trained ResNet-18 model
# For older torchvision versions, you might use pretrained=True
# model = models.resnet18(pretrained=True)
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model = model.to(device)
model.eval() # Set the model to evaluation mode

print("Original ResNet-18 model structure (first few layers and classifier):")
# print(model) # This will print the entire model structure

Setting model.eval() is important because it deactivates layers like Dropout and Batch Normalization’s training behavior, ensuring consistent feature extraction.

Step 3: Prepare an Input Image

Pre-trained models expect input images to be preprocessed in a specific way (e.g., resized, normalized). The weights object often provides the recommended transforms.

# Get the recommended transforms for the ResNet-18 model
transforms_pipeline = models.ResNet18_Weights.DEFAULT.transforms()
print("\nImage transforms pipeline:")
print(transforms_pipeline)

# Create a dummy image or load a real one
# For a real image:
# try:
#   img = Image.open("your_image.jpg").convert("RGB")
# except FileNotFoundError:
#   print("Dummy image will be used as 'your_image.jpg' was not found.")
#   img = Image.new('RGB', (224, 224), color = 'red') # Dummy image

# For demonstration, let's create a random tensor representing a batch of one image
dummy_image_tensor = torch.randn(1, 3, 224, 224).to(device) # Batch_size=1, 3 channels, 224x224 pixels

The input tensor should typically be of shape (batch_size, channels, height, width).

Method 1: Using Model Children (Accessing Layers Directly)

If you want features from a layer just before the final classification layer, you can often modify the model by replacing its classifier part.

# ResNet-18's final fully connected layer is named 'fc'
# We can remove it or replace it with an identity layer to get features from the layer before it.
feature_extractor_v1 = nn.Sequential(*list(model.children())[:-1]) # Remove the last layer (fc)
feature_extractor_v1 = feature_extractor_v1.to(device)
feature_extractor_v1.eval()

with torch.no_grad(): # Disable gradient calculations for inference
    features_v1 = feature_extractor_v1(dummy_image_tensor)

print(f"\nFeatures from v1 (before fc layer, after avgpool): {features_v1.shape}")
# Expected output shape: torch.Size([1, 512, 1, 1]) for ResNet-18 (after adaptive avg pooling)
# We can flatten this: features_v1.squeeze().shape would be torch.Size([512]) for a single image

Method 2: Using Forward Hooks

Forward hooks are a flexible way to get activations from any intermediate layer without modifying the model’s architecture. A hook is a function that gets registered to a module and is executed during the forward pass.

extracted_features = {}

def get_features_hook(module_name):
    """Returns a hook function that stores the output of a module."""
    def hook(module, input, output):
        # For convolutional layers, output is a 4D tensor (batch, channels, height, width)
        # For linear layers, output is typically a 2D tensor (batch, features)
        extracted_features[module_name] = output.detach() # Detach to prevent gradient tracking
    return hook

# Example: Extract features from the output of the 'layer3' block in ResNet-18
# You need to know the name of the layer. You can find this by printing model architecture.
# print(model)

# For ResNet-18, 'layer3' is a Sequential block. We might want features after its last sub-block.
# Let's target model.layer3[-1] (the last Bottleneck/BasicBlock in layer3)
target_layer_name = 'layer3_last_block'
handle = model.layer3[-1].register_forward_hook(get_features_hook(target_layer_name))

# Perform a forward pass to trigger the hook
with torch.no_grad():
    _ = model(dummy_image_tensor)

# Remove the hook after use if you don't need it anymore
handle.remove()

if target_layer_name in extracted_features:
    features_v2 = extracted_features[target_layer_name]
    print(f"\nFeatures from v2 (hook on '{target_layer_name}'): {features_v2.shape}")
else:
    print(f"\nCould not extract features for '{target_layer_name}'. Check layer name and model structure.")

Understanding the Code

• model.eval(): Crucial for ensuring layers like BatchNorm and Dropout behave correctly during inference.
• nn.Sequential(*list(model.children())[:-1]): This creates a new model comprising all children of the original model except the last one. This is a common way to get features from the penultimate layer.
• Forward Hooks (register_forward_hook): A function is registered to a specific module (layer). When data passes through this module during the forward pass, the hook function is called with the module itself, its input, and its output. This allows capturing intermediate activations.
• output.detach(): Used within the hook to get a tensor that doesn’t track gradients, which is important for feature extraction when you’re not training.
• torch.no_grad(): Context manager that disables gradient calculation, reducing memory consumption and speeding up inference.

Practical Applications

Extracted features can be used for:

• Transfer Learning: Use the features as input to a new, smaller model for a different task, especially when you have limited data for the new task.
• Image Retrieval/Similarity: Compare feature vectors of different images to find similar ones.
• Visualization: Techniques like t-SNE or UMAP can be applied to feature vectors to visualize how the model clusters different types of images.

Choosing which layer to extract features from depends on the task. Earlier layers capture more general, low-level features (edges, textures), while deeper layers capture more abstract, high-level features specific to the objects the model was trained on.