feat: 用 Grad-CAM 从分类模型生成弱监督框标签

分类模型只学过"是什么缺陷",Grad-CAM 把分类分数对卷积特征的梯度当作通道重要性,
加权求和得到空间热力图,阈值化 + 连通域外接矩形即得到框标签,输出 YOLO 格式,
可直接接入既有的检测训练链路。

- engine/gradcam.py:整批一次 backward 的 Grad-CAM(各样本只回传自己那一类的分数)、
  numpy 实现的 4-邻接连通域标记(不为这一处引入 scipy/opencv)、CAM 转归一化框
- scripts/gradcam_boxes.py:遍历数据集生成标签,并输出「原图 | 热力图 | 叠加+框」
  三联图供人工核对;框的类别取自图片所在目录而非模型预测

两处按实图定下的默认值:
- --stage 默认 6(layer3, 8×8)。layer4 只有 4×4 太粗,layer2 响应太散导致 40% 出不了框
- --peak-ratio 0.75:次要连通域峰值需达全局最大值的 75%,否则会在无缺陷的角落刷假框;
  实测把两框比例从 53% 降到 20%

全量 11,212 张已生成:61.8% 出 1 框、35.2% 出 2 框、3.0% 无框,每类留 25 张可视化。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-10 15:46:30 +08:00
co-authored by Claude Opus 5
parent 33fdf396ee
commit fc7136fddf
4 changed files with 631 additions and 0 deletions
+45
View File
@@ -632,6 +632,51 @@ model:
只加载 backbone,分类头仍随机初始化。`frozen_backbone: true` 可以单独衡量
自监督特征本身有多好。
## 从分类到定位:Grad-CAM 生成框标签
分类模型只学过"这是什么缺陷",没学过"缺陷在哪"。Grad-CAM 把分类分数对卷积特征的梯度
当作通道重要性,加权求和得到空间热力图——这是分类模型隐含的定位线索。阈值化 +
取连通域外接矩形,就得到弱监督的框标签:
```bash
uv run python scripts/gradcam_boxes.py --checkpoint runs/classify/best.pth
```
```
runs/gradcam_boxes/
├── labels/<类别>/<图名>.txt # YOLO 格式 class_id cx cy w h,可直接喂给检测训练
├── viz/<类别>/<图名>.png # 原图 | 热力图 | 叠加+框,供人工核对
├── data.yaml
└── summary.json
```
**框的类别取自图片所在目录**(人工归类结果),不是模型预测——即使模型分错,
要的也是该图真实类别的热力区域。
### 分辨率是主要限制
resnet18 在 128×128 输入下各层的 CAM 尺寸:layer4 只有 4×4、layer3 是 8×8、layer2 是 16×16。
层越浅越精细,但类别判别性越弱。`--stage`(4~7 对应 layer1~4)可选,默认 6 即 **layer3**:
| stage | CAM | 实测(每类抽 3 张,共 45 张) |
| --- | --- | --- |
| 7 (layer4) | 4×4 | 100% 出 1 个框,但只有 16 个格子,框很粗 |
| **6 (layer3)** | **8×8** | **76% 出 1 框、20% 出 2 框,热力区域与缺陷贴合** |
| 5 (layer2) | 16×16 | 40% 出不了框,响应太散 |
### 几个调参点
- `--threshold`(默认 0.4):热力图相对最大值的阈值
- `--peak-ratio`(默认 0.75):次要连通域的峰值要达到全局最大值的这个比例才保留。
不加这条会在没有缺陷的角落刷出假框——实测把两框比例从 53% 降到 20%
- `--padding`(默认 0.1):框四周外扩。CAM 倾向只覆盖最具判别性的部位(缺陷中心),
往往比真实缺陷范围小
- `--viz-per-class`(默认 20):每类存多少张可视化;`--viz-all` 全存(很慢很占盘)
> **不是所有类别都可定位。** `normal`/`black`/`edge_clean` 这类描述的是整图属性而非
> 局部缺陷,它们的框没有物理意义。真正有用的是 `blackdot`/`whitedot`/`mosquito`/
> `grease`/`zigzag` 这些有局部目标的类别。
## 方法
- **骨架**:`tiny_conv`(自研深度可分离卷积,默认)或 `torchvision`(resnet / mobilenet / shufflenet)。
+159
View File
@@ -0,0 +1,159 @@
"""Grad-CAM 与由热力图导出检测框。
分类模型只学过"这是什么缺陷",没学过"缺陷在哪"。Grad-CAM 把分类分数对某层特征图的
梯度当作该层各通道的重要性,加权求和后得到一张空间热力图——分类模型隐含的定位线索。
再对热力图阈值化、取连通域外接矩形,就得到弱监督的框标签。
分辨率是主要限制:resnet18 在 128×128 输入下 layer4 只有 4×4、layer3 有 8×8。
层越浅空间越精细,但类别判别性越弱,需要按数据实测取舍。
"""
from __future__ import annotations
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
def grad_cam(
backbone: nn.Module,
head: nn.Module,
images: torch.Tensor,
targets: torch.Tensor,
stage: int | None = None,
) -> torch.Tensor:
"""对一个 batch 求 Grad-CAM,返回 (B, h, w) 的热力图(每张已归一化到 [0, 1])。
Args:
stage: 取 `backbone.features` 的第几层输出做 CAM;None 表示用骨架最终输出
(即最后一个卷积块)。层越浅分辨率越高、类别判别性越弱。
对每个样本只回传它自己那一类的分数,所以整批可以一次 backward——
样本之间的梯度互不干扰。
"""
layers = getattr(backbone, "features", None)
if stage is not None and not isinstance(layers, nn.Sequential):
raise ValueError("该骨架不是 nn.Sequential,无法按 stage 取中间层")
with torch.enable_grad():
if stage is None:
activations = backbone(images)
features = activations
else:
features = images
for index, layer in enumerate(layers):
features = layer(features)
if index == stage:
break
else:
raise ValueError(f"stage {stage} 超出骨架层数 {len(layers)}")
activations = features
for layer in list(layers)[stage + 1 :]:
features = layer(features)
activations.retain_grad()
logits = head(features)
# 只取每个样本对应目标类别的分数;求和后一次 backward 即可
score = logits.gather(1, targets.view(-1, 1)).sum()
backbone.zero_grad(set_to_none=True)
head.zero_grad(set_to_none=True)
score.backward()
gradients = activations.grad
if gradients is None:
raise RuntimeError("没有拿到梯度,检查目标层是否参与了前向计算")
weights = gradients.mean(dim=(2, 3), keepdim=True) # 各通道的重要性
cam = F.relu((weights * activations).sum(dim=1)) # (B, h, w)
flat = cam.flatten(1)
lo = flat.min(dim=1, keepdim=True).values
hi = flat.max(dim=1, keepdim=True).values
cam = ((flat - lo) / (hi - lo).clamp(min=1e-8)).view_as(cam)
return cam.detach()
def connected_components(mask: np.ndarray) -> list[np.ndarray]:
"""4-邻接连通域标记,返回每个连通域的布尔掩码。
自己写而不引 scipy/opencv:项目只依赖 numpy/torch/PIL,
为这一处再加一个重依赖不划算。
"""
height, width = mask.shape
seen = np.zeros_like(mask, dtype=bool)
components = []
for start_y in range(height):
for start_x in range(width):
if not mask[start_y, start_x] or seen[start_y, start_x]:
continue
stack = [(start_y, start_x)]
seen[start_y, start_x] = True
pixels = []
while stack:
y, x = stack.pop()
pixels.append((y, x))
for ny, nx in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
if 0 <= ny < height and 0 <= nx < width and mask[ny, nx] and not seen[ny, nx]:
seen[ny, nx] = True
stack.append((ny, nx))
component = np.zeros_like(mask, dtype=bool)
ys, xs = zip(*pixels, strict=True)
component[list(ys), list(xs)] = True
components.append(component)
return components
def cam_to_boxes(
cam: np.ndarray,
threshold: float = 0.4,
min_area_ratio: float = 0.01,
max_boxes: int = 3,
padding: float = 0.0,
peak_ratio: float = 0.75,
) -> list[tuple[float, float, float, float]]:
"""把一张热力图变成若干归一化 xyxy 框。
Args:
threshold: 相对最大值的阈值;CAM 已归一化到 [0,1],所以直接比较。
min_area_ratio: 连通域面积小于整图这个比例就丢掉,滤掉零星的高响应点。
max_boxes: 最多保留几个框(按连通域内热力总和排序)。
padding: 框的四周外扩比例——CAM 倾向只覆盖最具判别性的部位(比如缺陷中心),
往往比真实缺陷范围小,适当外扩更贴合。
peak_ratio: 连通域的峰值至少要达到全局最大值的这个比例才保留。
CAM 归一化后最强的区域峰值恒为 1,次强区域若只有 0.45 多半是噪声——
实测不加这条会在没有缺陷的角落刷出假框。
返回的坐标已归一化到 [0, 1],与图像实际尺寸无关,可直接写 YOLO 标签。
"""
height, width = cam.shape
mask = cam >= threshold
if not mask.any():
return []
boxes = []
for component in connected_components(mask):
if component.sum() < max(1, min_area_ratio * height * width):
continue
if cam[component].max() < peak_ratio: # 弱响应区多半是噪声,不是第二个缺陷
continue
ys, xs = np.where(component)
# +1 让框覆盖住整个格子,而不是格子中心之间的距离
x1, x2 = xs.min() / width, (xs.max() + 1) / width
y1, y2 = ys.min() / height, (ys.max() + 1) / height
if padding:
dx, dy = (x2 - x1) * padding, (y2 - y1) * padding
x1, x2 = max(0.0, x1 - dx), min(1.0, x2 + dx)
y1, y2 = max(0.0, y1 - dy), min(1.0, y2 + dy)
boxes.append((float(cam[component].sum()), (x1, y1, x2, y2)))
boxes.sort(key=lambda item: -item[0])
return [box for _, box in boxes[:max_boxes]]
def to_yolo_line(box: tuple[float, float, float, float], class_id: int) -> str:
"""归一化 xyxy -> YOLO 的 `class_id cx cy w h`。"""
x1, y1, x2, y2 = box
return (
f"{class_id} {(x1 + x2) / 2:.6f} {(y1 + y2) / 2:.6f} {x2 - x1:.6f} {y2 - y1:.6f}"
)
+247
View File
@@ -0,0 +1,247 @@
#!/usr/bin/env python
"""用训练好的分类模型 + Grad-CAM 为每张图生成框标签。
uv run python scripts/gradcam_boxes.py --checkpoint runs/classify/best.pth
分类模型只学过"这是什么缺陷"。Grad-CAM 把分类分数对卷积特征的梯度当作通道重要性,
加权求和得到空间热力图——这是分类模型隐含的定位线索。阈值化 + 取连通域外接矩形,
就得到弱监督的框标签。
产出:
```
runs/gradcam_boxes/
├── labels/<类别>/<图名>.txt # YOLO 格式:class_id cx cy w h(归一化)
├── viz/<类别>/<图名>.png # 原图 | 热力图 | 叠加+框,供人工核对匹配程度
├── data.yaml # 类别表,可直接喂给检测训练
└── summary.json
```
**框的类别取自图片所在目录**(人工归类的结果),不是模型预测——即使模型分错,
我们要的也是该图真实类别的热力区域。
分辨率是主要限制:resnet18 在 128×128 输入下 layer4 只有 4×4、layer3 有 8×8。
`--stage` 可选层(4~7 对应 layer1~4),层越浅越精细但类别判别性越弱。
"""
from __future__ import annotations
import argparse
import json
import sys
from collections import Counter
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np # noqa: E402
import torch # noqa: E402
from PIL import Image # noqa: E402
from torch.utils.data import DataLoader # noqa: E402
from datasets import FolderClassificationDataset, classification_collate # noqa: E402
from datasets.yolo import write_data_yaml # noqa: E402
from engine import load_checkpoint # noqa: E402
from engine.gradcam import cam_to_boxes, grad_cam, to_yolo_line # noqa: E402
from models import build_model # noqa: E402
from utils import load_config, setup_logging # noqa: E402
from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402
logger = None
def save_visualization(path: Path, image_path: Path, cam: np.ndarray, boxes, title: str) -> None:
"""三联图:原图 | 热力图 | 叠加+框。人工一眼就能看出热力区域与框合不合。"""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
with Image.open(image_path) as raw:
image = np.asarray(raw.convert("RGB"))
height, width = image.shape[:2]
heat = np.asarray(
Image.fromarray((cam * 255).astype(np.uint8)).resize((width, height), Image.BILINEAR)
) / 255.0
fig, axes = plt.subplots(1, 3, figsize=(11, 4))
axes[0].imshow(image)
axes[0].set_title("original")
axes[1].imshow(heat, cmap="jet", vmin=0, vmax=1)
axes[1].set_title(f"grad-cam {cam.shape[0]}x{cam.shape[1]}")
axes[2].imshow(image)
axes[2].imshow(heat, cmap="jet", alpha=0.45, vmin=0, vmax=1)
for x1, y1, x2, y2 in boxes:
axes[2].add_patch(
Rectangle(
(x1 * width, y1 * height), (x2 - x1) * width, (y2 - y1) * height,
fill=False, edgecolor="lime", linewidth=2,
)
)
axes[2].set_title(f"overlay + {len(boxes)} box")
for ax in axes:
ax.axis("off")
fig.suptitle(title, fontsize=10)
fig.tight_layout()
path.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(path, dpi=110)
plt.close(fig)
def main() -> int:
global logger
parser = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
parser.add_argument("--checkpoint", default="runs/classify/best.pth", help="分类模型权重")
parser.add_argument("--config", default="configs/classify/film.yaml", help="训练时的配置")
parser.add_argument("--out", default="runs/gradcam_boxes", help="输出目录")
parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
parser.add_argument("--device", default=None)
parser.add_argument(
"--stage", type=int, default=6,
help="取骨架 features 的第几层做 CAM(4~7 = layer1~4);默认 6 即 layer3",
)
parser.add_argument("--threshold", type=float, default=0.4, help="热力图阈值(相对最大值)")
parser.add_argument("--padding", type=float, default=0.1, help="框四周外扩比例")
parser.add_argument("--max-boxes", type=int, default=2, help="每张图最多几个框")
parser.add_argument(
"--min-area-ratio", type=float, default=0.02, help="连通域面积下限(占整图比例)"
)
parser.add_argument(
"--peak-ratio", type=float, default=0.75,
help="次要连通域的峰值至少要达到全局最大值的这个比例才保留,用于滤掉弱响应假框",
)
parser.add_argument(
"--viz-per-class", type=int, default=20, help="每个类别存多少张可视化,0 = 不存"
)
parser.add_argument("--viz-all", action="store_true", help="每张图都存可视化(很慢很占盘)")
parser.add_argument("--limit", type=int, default=0, help="只处理前 N 张,用于快速试参数")
args = parser.parse_args()
out_dir = Path(args.out)
out_dir.mkdir(parents=True, exist_ok=True)
logger = setup_logging(out_dir / "gradcam.log")
cfg = load_config(args.config, args.set)
try:
device = resolve_device(args.device or cfg.get("device", "auto"))
except DeviceNotAvailableError as exc:
logger.error("%s", exc)
return 1
data_cfg = cfg["data"]
dataset = FolderClassificationDataset(
root=data_cfg["root"],
split="all", # 为每个样本都生成,不分训练/验证
classes=data_cfg.get("classes"),
transforms=data_cfg["val_transforms"],
exclude=data_cfg.get("exclude") or [],
)
if args.limit:
dataset.samples = dataset.samples[: args.limit]
classes = dataset.classes
model_cfg = dict(cfg["model"])
model_cfg.setdefault("num_classes", len(classes))
model_cfg["init_from"] = None # 权重从 checkpoint 来,不要再读自监督骨架
model = build_model(model_cfg)
load_checkpoint(args.checkpoint, model, map_location=device)
model.to(device).eval()
for param in model.parameters(): # Grad-CAM 需要梯度,但不更新权重
param.requires_grad_(True)
logger.info("设备: %s | 权重: %s", describe_device(device), args.checkpoint)
logger.info("数据集 %s:%d 张,%d 类", data_cfg["root"], len(dataset), len(classes))
logger.info(
"CAM 目标层 features[%d] | 阈值 %.2f | 外扩 %.2f | 每图最多 %d 框",
args.stage, args.threshold, args.padding, args.max_boxes,
)
loader = DataLoader(
dataset,
batch_size=int(data_cfg.get("loader", {}).get("batch_size", 64)),
shuffle=False,
num_workers=int(data_cfg.get("loader", {}).get("num_workers", 4)),
collate_fn=classification_collate,
)
labels_dir = out_dir / "labels"
viz_dir = out_dir / "viz"
viz_done: Counter[str] = Counter()
box_counts: Counter[int] = Counter()
per_class_boxes: Counter[str] = Counter()
correct = 0
cam_shape = None
for step, batch in enumerate(loader):
images = batch["images"].to(device)
targets = batch["labels"].to(device)
cams = grad_cam(model.backbone, model.head, images, targets, stage=args.stage)
with torch.no_grad():
correct += int((model(images).argmax(1) == targets).sum())
cam_shape = tuple(cams.shape[-2:])
for cam, target, meta in zip(cams.cpu().numpy(), targets.cpu().tolist(), batch["meta"],
strict=True):
name = classes[target]
boxes = cam_to_boxes(
cam,
threshold=args.threshold,
min_area_ratio=args.min_area_ratio,
max_boxes=args.max_boxes,
padding=args.padding,
peak_ratio=args.peak_ratio,
)
box_counts[len(boxes)] += 1
per_class_boxes[name] += len(boxes)
source = Path(meta["path"])
label_path = labels_dir / name / f"{source.stem}.txt"
label_path.parent.mkdir(parents=True, exist_ok=True)
label_path.write_text(
"\n".join(to_yolo_line(box, target) for box in boxes) + ("\n" if boxes else ""),
encoding="utf-8",
)
if args.viz_all or viz_done[name] < args.viz_per_class:
viz_done[name] += 1
save_visualization(
viz_dir / name / f"{source.stem}.png", source, cam, boxes,
f"{name} | {len(boxes)} box | thr={args.threshold}",
)
if step % 20 == 0:
logger.info("处理 [%d/%d]", step, len(loader))
total = sum(box_counts.values())
write_data_yaml(out_dir / "data.yaml", classes, root=out_dir)
summary = {
"checkpoint": args.checkpoint,
"dataset": str(data_cfg["root"]),
"num_images": total,
"cam_shape": list(cam_shape) if cam_shape else None,
"stage": args.stage,
"threshold": args.threshold,
"padding": args.padding,
"max_boxes": args.max_boxes,
"min_area_ratio": args.min_area_ratio,
"peak_ratio": args.peak_ratio,
"classification_accuracy_on_all": correct / total if total else 0.0,
"boxes_per_image": {str(k): v for k, v in sorted(box_counts.items())},
"boxes_per_class": dict(sorted(per_class_boxes.items())),
}
(out_dir / "summary.json").write_text(
json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
)
logger.info("CAM 分辨率 %s | 全量分类准确率 %.4f", cam_shape, summary["classification_accuracy_on_all"])
for count, number in sorted(box_counts.items()):
logger.info(" 生成 %d 个框的图片: %d 张 (%.1f%%)", count, number, 100 * number / total)
logger.info("标签写入 %s,可视化 %d 张写入 %s", labels_dir, sum(viz_done.values()), viz_dir)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+180
View File
@@ -0,0 +1,180 @@
"""Grad-CAM 与由热力图导出框。"""
import pytest
torch = pytest.importorskip("torch")
import numpy as np # noqa: E402
from torch import nn # noqa: E402
from engine.gradcam import ( # noqa: E402
cam_to_boxes,
connected_components,
grad_cam,
to_yolo_line,
)
from models import build_model # noqa: E402
def test_connected_components_separates_blobs():
mask = np.array([[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 1]], dtype=bool)
components = connected_components(mask)
assert len(components) == 2
assert sorted(int(c.sum()) for c in components) == [3, 4]
assert not (components[0] & components[1]).any() # 互不重叠
assert (components[0] | components[1] == mask).all()
def test_connected_components_is_four_connected():
"""对角相邻不算连通——否则细长划痕会和邻近噪点粘成一个框。"""
mask = np.array([[1, 0], [0, 1]], dtype=bool)
assert len(connected_components(mask)) == 2
def test_cam_to_boxes_finds_the_hot_region():
cam = np.zeros((8, 8), dtype=np.float32)
cam[2:5, 3:7] = 1.0
(box,) = cam_to_boxes(cam, threshold=0.5, min_area_ratio=0.0, peak_ratio=0.75)
x1, y1, x2, y2 = box
assert (x1, x2) == pytest.approx((3 / 8, 7 / 8)) # 框覆盖整格,不是格心距离
assert (y1, y2) == pytest.approx((2 / 8, 5 / 8))
def test_padding_expands_within_bounds():
cam = np.zeros((8, 8), dtype=np.float32)
cam[3:5, 3:5] = 1.0
plain = cam_to_boxes(cam, threshold=0.5, min_area_ratio=0.0)[0]
padded = cam_to_boxes(cam, threshold=0.5, min_area_ratio=0.0, padding=0.5)[0]
assert padded[0] < plain[0] and padded[2] > plain[2]
assert all(0.0 <= v <= 1.0 for v in padded) # 外扩不会越界
corner = np.zeros((8, 8), dtype=np.float32)
corner[0:2, 0:2] = 1.0
assert cam_to_boxes(corner, threshold=0.5, min_area_ratio=0.0, padding=1.0)[0][:2] == (0.0, 0.0)
def test_peak_ratio_drops_weak_secondary_blobs():
"""次强区域峰值只有 0.45 时多半是噪声——实测不过滤会在无缺陷处刷出假框。"""
cam = np.zeros((8, 8), dtype=np.float32)
cam[1:3, 1:3] = 1.0 # 主区域
cam[5:7, 5:7] = 0.45 # 弱响应区
assert len(cam_to_boxes(cam, threshold=0.4, min_area_ratio=0.0, peak_ratio=0.0)) == 2
assert len(cam_to_boxes(cam, threshold=0.4, min_area_ratio=0.0, peak_ratio=0.75)) == 1
def test_min_area_and_max_boxes_limits():
cam = np.zeros((10, 10), dtype=np.float32)
cam[0:1, 0:1] = 1.0 # 1 格,面积很小
cam[4:8, 4:8] = 1.0
assert len(cam_to_boxes(cam, threshold=0.5, min_area_ratio=0.05, peak_ratio=0.5)) == 1
many = np.zeros((12, 12), dtype=np.float32)
for offset in (0, 4, 8):
many[offset:offset + 2, offset:offset + 2] = 1.0
assert len(cam_to_boxes(many, threshold=0.5, min_area_ratio=0.0, max_boxes=2)) == 2
def test_no_box_when_nothing_activates():
assert cam_to_boxes(np.zeros((8, 8), dtype=np.float32), threshold=0.5) == []
def test_to_yolo_line_roundtrip():
line = to_yolo_line((0.2, 0.4, 0.6, 0.8), class_id=7)
parts = line.split()
assert parts[0] == "7"
cx, cy, w, h = (float(v) for v in parts[1:])
assert (cx, cy, w, h) == pytest.approx((0.4, 0.6, 0.4, 0.4))
class _TinyNet(nn.Module):
"""features 是 nn.Sequential,好按 stage 取中间层。"""
def __init__(self) -> None:
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 4, 3, padding=1), nn.ReLU(),
nn.Conv2d(4, 8, 3, stride=2, padding=1), nn.ReLU(),
)
self.out_channels = 8
def forward(self, x):
return self.features(x)
def test_grad_cam_shapes_and_normalization():
torch.manual_seed(0)
backbone = _TinyNet()
head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 3))
images = torch.randn(2, 3, 16, 16)
targets = torch.tensor([0, 2])
cam = grad_cam(backbone, head, images, targets)
assert cam.shape == (2, 8, 8) # 与最终特征图同分辨率
assert not cam.requires_grad
assert float(cam.min()) >= 0.0 and float(cam.max()) <= 1.0
# 每张单独归一化:非退化的那张峰值应当正好是 1
peaks = cam.flatten(1).max(dim=1).values
assert any(float(p) == pytest.approx(1.0) for p in peaks)
def test_grad_cam_all_zero_when_relu_kills_everything():
"""梯度全为负时 ReLU 会清零整张图,此时归一化结果仍是合法的全 0,
对应下游"这张图提取不出框"——全量数据里约 3% 属于这种情况。"""
backbone = _TinyNet()
head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 3))
with torch.no_grad(): # 让目标类别的权重全为负,梯度随之为负
head[2].weight.fill_(-1.0)
head[2].bias.zero_()
with torch.no_grad():
backbone.features[0].weight.fill_(0.1)
backbone.features[2].weight.fill_(0.1)
cam = grad_cam(backbone, head, torch.rand(1, 3, 16, 16), torch.tensor([0]))
assert float(cam.max()) == pytest.approx(0.0)
assert cam_to_boxes(cam[0].numpy(), threshold=0.4) == []
def test_grad_cam_stage_controls_resolution():
backbone = _TinyNet()
head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 3))
images = torch.randn(1, 3, 16, 16)
targets = torch.tensor([1])
# 浅层分辨率更高:stage 1 在第一次下采样之前
assert grad_cam(backbone, head, images, targets, stage=1).shape == (1, 16, 16)
assert grad_cam(backbone, head, images, targets, stage=3).shape == (1, 8, 8)
with pytest.raises(ValueError, match="超出骨架层数"):
grad_cam(backbone, head, images, targets, stage=99)
def test_grad_cam_responds_to_the_target_class():
"""不同目标类别应当给出不同的热力图,否则 CAM 就没有类别判别性。"""
torch.manual_seed(0)
backbone = _TinyNet()
head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 3))
images = torch.randn(1, 3, 16, 16)
a = grad_cam(backbone, head, images, torch.tensor([0]))
b = grad_cam(backbone, head, images, torch.tensor([1]))
assert not torch.allclose(a, b)
def test_grad_cam_works_on_the_real_classifier():
model = build_model({
"name": "classifier", "num_classes": 5,
"backbone": {"name": "torchvision", "arch": "resnet18", "pretrained": False},
}).eval()
for param in model.parameters():
param.requires_grad_(True)
images = torch.randn(2, 3, 128, 128)
targets = torch.tensor([1, 4])
cam = grad_cam(model.backbone, model.head, images, targets, stage=6)
assert cam.shape == (2, 8, 8) # layer3 在 128 输入下是 8x8
boxes = cam_to_boxes(cam[0].cpu().numpy(), threshold=0.4, min_area_ratio=0.0)
assert all(0.0 <= v <= 1.0 for box in boxes for v in box)