| 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() |
