Create model.py

model.py

@@ -0,0 +1,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()