diff --git a/README.md b/README.md index 3308677..ba4251d 100644 --- a/README.md +++ b/README.md @@ -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)。 diff --git a/engine/gradcam.py b/engine/gradcam.py new file mode 100644 index 0000000..bbcc153 --- /dev/null +++ b/engine/gradcam.py @@ -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}" + ) diff --git a/scripts/gradcam_boxes.py b/scripts/gradcam_boxes.py new file mode 100644 index 0000000..db233db --- /dev/null +++ b/scripts/gradcam_boxes.py @@ -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()) diff --git a/tests/test_gradcam.py b/tests/test_gradcam.py new file mode 100644 index 0000000..2f7ca21 --- /dev/null +++ b/tests/test_gradcam.py @@ -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)