Lesson 38: Visualizing and Understanding CNNs¶
A trained CNN is a black box in the sense that its millions of weights don't have obvious individual meanings. But where in the input image a prediction comes from is answerable, and answering it is often what separates "the model got the right answer" from "the model got the right answer for the right reason." This lesson builds two visualization tools from scratch: saliency maps (Simonyan et al., 2013) and Grad-CAM (Selvaraju et al., 2017).
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
Setup: the plus-vs-circle CNN from Lesson 33, with feature maps exposed¶
The only change from Lesson 33's CNNClassifier is that forward now also returns the last convolutional layer's feature map (before global pooling), so both visualization methods below can get at it.
def make_image(shape_type, cx, cy, size=16):
img = np.zeros((size, size), dtype=np.float32)
if shape_type == 'plus':
img[cy-1:cy+2, cx-3:cx+4] = 1.0
img[cy-3:cy+4, cx-1:cx+2] = 1.0
else:
yy, xx = np.mgrid[0:size, 0:size]
img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0
return img
def make_dataset(rng_local, n, position_range=(4, 12)):
imgs, labels, positions = [], [], []
for _ in range(n):
shape_type = rng_local.choice(['plus', 'circle'])
cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)
imgs.append(make_image(shape_type, cx, cy))
labels.append(0.0 if shape_type == 'plus' else 1.0)
positions.append((cx, cy))
return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.float32), positions
rng = np.random.default_rng(4)
X_train, y_train, _ = make_dataset(rng, 300)
X_test, y_test, pos_test = make_dataset(rng, 50)
class CNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 8, 5, padding=2)
self.conv2 = nn.Conv2d(8, 16, 5, padding=2)
self.gpool = nn.AdaptiveMaxPool2d(1)
self.fc = nn.Linear(16, 1)
def forward(self, x):
f1 = F.relu(self.conv1(x))
f2 = F.relu(self.conv2(f1)) # last conv feature map, full 16x16 resolution
feat = self.gpool(f2).flatten(1)
return self.fc(feat).squeeze(-1), f2
torch.manual_seed(0)
model = CNN()
opt = torch.optim.Adam(model.parameters(), lr=0.01)
Xt = torch.tensor(X_train).unsqueeze(1); yt = torch.tensor(y_train)
for _ in range(300):
opt.zero_grad()
out, _ = model(Xt)
loss = F.binary_cross_entropy_with_logits(out, yt)
loss.backward()
opt.step()
with torch.no_grad():
out, _ = model(torch.tensor(X_test).unsqueeze(1))
acc = ((out > 0).float() == torch.tensor(y_test)).float().mean().item()
print(f'test accuracy: {acc:.1%}')
Saliency maps¶
The idea (Simonyan et al., 2013): take the gradient of the predicted class score with respect to every input pixel. A pixel with a large-magnitude gradient is one where a small change would most change the prediction — i.e., a pixel the network is "looking at."
def saliency_map(model, img):
x = torch.tensor(img[None, None]).float()
x.requires_grad_(True)
score, feat = model(x)
score.backward()
return x.grad[0, 0].abs().numpy(), feat
idx = 3
saliency, _ = saliency_map(model, X_test[idx])
cx, cy = pos_test[idx]
fig, axes = plt.subplots(1, 2, figsize=(7, 3.2))
axes[0].imshow(X_test[idx], cmap='gray')
axes[0].scatter([cx], [cy], c='red', marker='x', s=60, label='true center')
axes[0].set_title('input image'); axes[0].legend(fontsize=7); axes[0].axis('off')
axes[1].imshow(saliency, cmap='hot')
axes[1].set_title('saliency map')
axes[1].axis('off')
plt.show()
Grad-CAM¶
Raw saliency maps are pixel-level and tend to be noisy. Grad-CAM (Selvaraju et al., 2017) instead works on the last convolutional layer's feature maps, which are lower-resolution but far more semantically meaningful:
- Take the gradient of the class score with respect to each channel of the last conv feature map.
- Average each channel's gradient spatially to get one importance weight per channel.
- Form a weighted sum of the feature-map channels using those weights, then apply ReLU (only positive evidence for the class matters).
The result is a coarse heatmap, the same spatial size as the last conv layer, that can be upsampled back to the input resolution.
def grad_cam(model, img, out_size=16):
x = torch.tensor(img[None, None]).float()
x.requires_grad_(True)
score, feat = model(x)
feat.retain_grad()
score.backward()
weights = feat.grad[0].mean(dim=(1, 2)) # (channels,) importance per channel
cam = F.relu((weights[:, None, None] * feat[0]).sum(dim=0))
cam_up = F.interpolate(cam[None, None], size=(out_size, out_size), mode='bilinear', align_corners=False)
return cam_up[0, 0].detach().numpy()
cam = grad_cam(model, X_test[idx])
fig, axes = plt.subplots(1, 3, figsize=(10, 3.2))
axes[0].imshow(X_test[idx], cmap='gray')
axes[0].scatter([cx], [cy], c='red', marker='x', s=60)
axes[0].set_title('input image'); axes[0].axis('off')
axes[1].imshow(saliency, cmap='hot')
axes[1].set_title('saliency map'); axes[1].axis('off')
axes[2].imshow(cam, cmap='hot')
axes[2].set_title('Grad-CAM'); axes[2].axis('off')
plt.show()
Do these maps actually point at the shape?¶
A single example is a nice picture but not evidence. Check quantitatively: for every test image, find each map's peak pixel and measure its distance to the shape's true center, and compare against the distance a random guess would get.
sal_dists, cam_dists, rand_dists = [], [], []
rand_rng = np.random.default_rng(0)
for i in range(len(X_test)):
sal, _ = saliency_map(model, X_test[i])
cam_i = grad_cam(model, X_test[i])
cx_i, cy_i = pos_test[i]
peak_sal = np.unravel_index(sal.argmax(), sal.shape) # (row, col) = (y, x)
peak_cam = np.unravel_index(cam_i.argmax(), cam_i.shape)
sal_dists.append(np.hypot(peak_sal[1] - cx_i, peak_sal[0] - cy_i))
cam_dists.append(np.hypot(peak_cam[1] - cx_i, peak_cam[0] - cy_i))
rx, ry = rand_rng.integers(0, 16), rand_rng.integers(0, 16)
rand_dists.append(np.hypot(rx - cx_i, ry - cy_i))
print(f'{"method":>18} {"mean peak distance to true center":>36}')
print(f'{"saliency map":>18} {np.mean(sal_dists):>33.2f} px')
print(f'{"Grad-CAM":>18} {np.mean(cam_dists):>33.2f} px')
print(f'{"random baseline":>18} {np.mean(rand_dists):>33.2f} px')
Both methods land far closer to the true shape center than a random guess would, confirming they're picking out genuinely relevant image regions rather than something spurious — and here, with full spatial resolution preserved in the last conv layer, Grad-CAM is actually the more precise of the two.
In practice these tools matter most for catching a specific failure: a network that gets the right answer on a training or validation set for the wrong reason — for example, learning to recognize a background watermark that happened to correlate with one class, rather than the object itself. A saliency map or Grad-CAM overlay on such a network would show it "looking" at the watermark, not the object — a bug that overall accuracy alone would never reveal (echoing Lesson 37's point that a single accuracy number hides a lot).
Exercise¶
- Grad-CAM here uses the last conv layer. Modify
grad_camto instead use the intermediate feature map afterconv1(beforeconv2). Does the resulting heatmap get sharper (closer to pixel-perfect, like the saliency map) or coarser, and why would an earlier layer behave that way? - Deliberately corrupt the dataset: add a small fixed bright square in the same corner of every "plus" training image (but no circle images), retrain, and check whether Grad-CAM on a test plus image lights up on the shape or on the corrupted corner.
- The saliency-map gradient in this lesson is taken with respect to the raw logit (
score), not the sigmoid probability. Try computing it with respect totorch.sigmoid(score)instead — does the resulting map look meaningfully different, and can you explain why using the chain rule?