File size: 2,298 Bytes
a36ad62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
class SceneCNN(nn.Module):
def __init__(self, base=16, num_classes=6, conv_dropout=0.1, fc_dropout=0.4):
super().__init__()
channels = [base, base * 2, base * 4, base * 8]
layers = OrderedDict()
in_ch = 3
for i, out_ch in enumerate(channels, 1):
layers[f"conv{i}a"] = nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False)
layers[f"bn{i}a"] = nn.BatchNorm2d(out_ch)
layers[f"relu{i}a"] = nn.ReLU()
layers[f"conv{i}b"] = nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False)
layers[f"bn{i}b"] = nn.BatchNorm2d(out_ch)
layers[f"relu{i}b"] = nn.ReLU()
layers[f"pool{i}"] = nn.MaxPool2d(2)
layers[f"drop{i}"] = nn.Dropout2d(conv_dropout)
in_ch = out_ch
self.features = nn.Sequential(layers)
self.gap = nn.AdaptiveAvgPool2d(1)
self.flatten = nn.Flatten()
self.drop_fc = nn.Dropout(fc_dropout)
self.fc1 = nn.Linear(channels[-1], 128)
self.relu_fc = nn.ReLU()
self.fc2 = nn.Linear(128, num_classes)
def forward(self, x):
x = self.features(x)
x = self.flatten(self.gap(x))
x = self.drop_fc(x)
x = self.relu_fc(self.fc1(x))
return self.fc2(x)
def grad_cam(model, x):
"""x: normalized tensor (3,H,W). Returns (cam HxW in [0,1], softmax probs)."""
model.eval()
store = {}
def hook(_, __, out):
store["act"] = out
out.register_hook(lambda g: store.__setitem__("grad", g))
h = model.features.relu4b.register_forward_hook(hook)
with torch.enable_grad():
logits = model(x.unsqueeze(0))
h.remove()
cls = logits.argmax(1).item()
model.zero_grad()
logits[0, cls].backward()
act, grad = store["act"][0], store["grad"][0]
weights = grad.mean(dim=(1, 2))
cam = F.relu((weights[:, None, None] * act).sum(0))
cam = F.interpolate(cam[None, None], size=x.shape[1:], mode="bilinear", align_corners=False)[0, 0]
cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
probs = F.softmax(logits, 1)[0]
return cam.detach().numpy(), probs.detach().numpy() |