Feature map visualization in PyTorch - ML Experiment: Train & Evaluate

Experiment - Feature map visualization

Problem:You have trained a convolutional neural network (CNN) on image data. You want to understand what features the network learns by visualizing the feature maps (outputs of convolutional layers) for a sample input image.

Current Metrics:Model trains with 85% accuracy on training data and 80% on validation data. No visualization of feature maps has been done yet.

Issue:Without visualizing feature maps, it is hard to interpret what the CNN focuses on in the images. This limits understanding and debugging of the model.

Your Task

Visualize the feature maps of the first convolutional layer for a given input image to understand what features the CNN extracts.

Use PyTorch framework.

Do not modify the trained model weights.

Visualize only the first convolutional layer's output feature maps.

Use matplotlib for plotting.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import matplotlib.pyplot as plt

# Define a simple CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, kernel_size=5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Load pretrained model (simulate training by loading random weights here)
model = SimpleCNN()
model.eval()

# Prepare a sample image from CIFAR10 validation set
transform = transforms.Compose([
    transforms.ToTensor(),
])

val_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=True)

# Get one sample image and label
sample_img, sample_label = next(iter(val_loader))

# Hook to capture feature maps
feature_maps = []

def hook_fn(module, input, output):
    feature_maps.append(output.detach())

# Register hook on first conv layer
hook = model.conv1.register_forward_hook(hook_fn)

# Forward pass
_ = model(sample_img)

# Remove hook
hook.remove()

# feature_maps[0] shape: [1, 6, 28, 28] (batch, channels, height, width)
fm = feature_maps[0][0]  # remove batch dimension

# Plot feature maps
num_maps = fm.shape[0]
cols = 3
rows = (num_maps + cols - 1) // cols
plt.figure(figsize=(cols * 3, rows * 3))
for i in range(num_maps):
    plt.subplot(rows, cols, i + 1)
    # Normalize to 0-1 for visualization
    fmap = fm[i]
    fmap = (fmap - fmap.min()) / (fmap.max() - fmap.min() + 1e-5)
    plt.imshow(fmap.cpu(), cmap='gray')
    plt.axis('off')
    plt.title(f'Feature map {i+1}')
plt.tight_layout()
plt.show()

Added a forward hook on the first convolutional layer to capture its output feature maps.

Extracted feature maps for a sample input image from the validation set.

Normalized each feature map to the 0-1 range for clear visualization.

Plotted all feature maps in a grid using matplotlib.

Added a small epsilon (1e-5) to denominator in normalization to avoid division by zero.

Results Interpretation

Before: No insight into what the CNN learns internally.

After: Visualized 6 feature maps from the first convolutional layer, each highlighting different edges and textures in the input image.

Visualizing feature maps helps understand what patterns the CNN detects early on. This aids in interpreting and debugging the model.

Bonus Experiment

Visualize feature maps from the second convolutional layer and compare them with the first layer's maps.

💡 Hint

Register a hook on the second conv layer and plot its feature maps similarly. Notice how deeper layers capture more complex features.

Practice

1. What is a feature map in the context of convolutional neural networks (CNNs)?

easy

A. The loss value during training

B. The input image to the CNN

C. The final prediction of the model

D. The output of a convolutional layer showing detected patterns

Solution

1. Step 1: Understand CNN layer outputs

   Convolutional layers process input images and produce outputs called feature maps that highlight detected features.

2. Step 2: Identify feature map role

   Feature maps represent learned patterns like edges or textures, not inputs or final outputs.

3. Final Answer:

   The output of a convolutional layer showing detected patterns -> Option D

4. Quick Check:

   Feature map = convolution output [OK]

Hint: Feature maps are outputs of conv layers, not inputs or predictions [OK]

Common Mistakes:

• Confusing feature maps with input images
• Thinking feature maps are final model outputs
• Mixing feature maps with loss values

2. Which PyTorch code snippet correctly extracts feature maps from a convolutional layer named conv1 given an input tensor x?

easy

A. feature_maps = x.conv1()

B. feature_maps = conv1(x)

C. feature_maps = conv1.forward()

D. feature_maps = conv1.output(x)

Solution

1. Step 1: Understand PyTorch layer call

   In PyTorch, calling a layer like a function with input tensor returns its output (feature maps).

2. Step 2: Check syntax correctness

   Using conv1(x) is correct; x.conv1() or conv1.output(x) are invalid syntax.

3. Final Answer:

   feature_maps = conv1(x) -> Option B

4. Quick Check:

   Call layer as function = correct [OK]

Hint: Call conv layer like a function with input tensor [OK]

Common Mistakes:

• Trying to call layer as method on input tensor
• Using non-existent methods like .output()
• Calling forward() directly instead of layer call

3. Given the following PyTorch code, what is the shape of feature_maps?

import torch
import torch.nn as nn
conv = nn.Conv2d(in_channels=3, out_channels=5, kernel_size=3, padding=1)
x = torch.randn(1, 3, 32, 32)
feature_maps = conv(x)

medium

A. [1, 3, 32, 32]

B. [1, 5, 30, 30]

C. [1, 5, 32, 32]

D. [3, 5, 32, 32]

Solution

1. Step 1: Analyze conv layer parameters

   Input has shape [1, 3, 32, 32]. Conv2d has 5 output channels, kernel size 3, padding 1.

2. Step 2: Calculate output spatial size

   Padding 1 keeps spatial size same: 32x32. Output channels = 5, batch size = 1.

3. Final Answer:

   [1, 5, 32, 32] -> Option C

4. Quick Check:

   Output shape = [batch, out_channels, height, width] [OK]

Hint: Padding keeps size; output channels define depth [OK]

Common Mistakes:

• Ignoring padding effect on output size
• Confusing input channels with output channels
• Mixing batch size with channel dimension

4. You try to visualize feature maps using this code but get an error:

import matplotlib.pyplot as plt
feature_maps = conv(x)
plt.imshow(feature_maps[0])
plt.show()

What is the likely cause of the error?

medium

A. feature_maps[0] has multiple channels, plt.imshow expects 2D or 3D image

B. conv(x) returns a scalar, not a tensor

C. plt.imshow cannot display tensors

D. x is not defined before conv(x)

Solution

1. Step 1: Understand feature_maps shape

   feature_maps[0] is shape [channels, height, width], multiple channels not a single image.

2. Step 2: plt.imshow expects 2D or 3D image

   plt.imshow needs 2D grayscale or 3D RGB image, but feature_maps[0] has multiple channels causing error.

3. Final Answer:

   feature_maps[0] has multiple channels, plt.imshow expects 2D or 3D image -> Option A

4. Quick Check:

   Multi-channel tensor ≠ single image [OK]

Hint: Plot one channel slice, not full multi-channel tensor [OK]

Common Mistakes:

• Trying to plot all channels at once with plt.imshow
• Assuming conv output is scalar
• Not checking input tensor existence

5. You want to visualize all feature maps from a convolutional layer conv for a single input image x. Which code correctly plots each channel as a separate grayscale image using matplotlib?

hard

A. feature_maps = conv(x)
for i in range(feature_maps.shape[1]):
plt.subplot(1, feature_maps.shape[1], i+1)
plt.imshow(feature_maps[0, i].detach().cpu(), cmap='gray')
plt.show()

B. feature_maps = conv(x)
plt.imshow(feature_maps.detach().cpu(), cmap='gray')
plt.show()

C. feature_maps = conv(x)
for i in range(feature_maps.shape[0]):
plt.imshow(feature_maps[i].detach().cpu(), cmap='gray')
plt.show()

D. feature_maps = conv(x)
plt.imshow(feature_maps[0].detach().cpu())
plt.show()

Solution

1. Step 1: Extract feature maps and iterate channels

   feature_maps shape is [batch, channels, height, width]. We select batch 0 and loop over channels.

2. Step 2: Plot each channel as grayscale image

   Use plt.subplot to arrange images, plt.imshow with cmap='gray' to show each channel properly.

3. Final Answer:

   feature_maps = conv(x) for i in range(feature_maps.shape[1]): plt.subplot(1, feature_maps.shape[1], i+1) plt.imshow(feature_maps[0, i].detach().cpu(), cmap='gray') plt.show() -> Option A

4. Quick Check:

   Loop channels, plot each with cmap='gray' [OK]

Hint: Loop channels, plot each with grayscale colormap [OK]

Common Mistakes:

• Plotting entire tensor at once
• Not detaching or moving tensor to CPU
• Ignoring batch dimension