refactor: 按实际的五步工作流裁掉无关代码
实际工作流是：simclr 训练 -> 聚类 -> 人工归类 -> 分类训练 -> GradCAM/SHAP 生成
热力图与 box。据此移除不参与其中的检测/定位链路，共 31 个文件、约 2600 行：
- 模型：defect_net、postprocess、necks/、heads/center_head
- 数据：film_defect（YOLO 读取端）、target_encoder
- 引擎：trainer、evaluator、losses
- 入口：train.py、eval.py
- 脚本：convert_annotations、split_dataset、export_onnx、profile_inference
- 配置：configs/{train,eval,model,dataset}
- utils/labels（多类/单类标签模式，仅检测分类头使用）
共享模块按实际用量瘦身：engine/metrics 只留 AverageMeter（检测指标随之移除）；
datasets/yolo 只留写出侧（IMAGE_SUFFIXES 与 write_data_yaml）——第 5 步产出的
YOLO 标签是给外部检测训练用的交换格式，读取端不再需要。
裁剪过程中修掉一个真 bug：generate_boxes 的 --stage 默认 6（layer3）对 tiny_conv
这类非 nn.Sequential 骨架直接抛异常。现在自动退回骨架最终输出并警告，
另加 --stage -1 显式选择最终输出。
五步工作流已端到端实跑验证（simclr 1 轮 -> 聚类 -> 导出 -> 分类 1 轮 ->
gradcam 与 shap 两种方法各生成一遍）。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
57446c2f10
commit
86c86b105c
+26
-102
@@ -1,96 +1,12 @@
|
||||
"""卷积骨架与分类模型。"""
|
||||
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
from models import build_model # noqa: E402
|
||||
from models.backbones.resnet import _SUPPORTED as _SUPPORTED_ARCHS # noqa: E402
|
||||
from models.postprocess import decode_heatmap # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_cfg():
|
||||
return {
|
||||
"name": "defect_net",
|
||||
"num_classes": 3,
|
||||
"backbone": {"name": "tiny_conv", "widths": [8, 16, 32], "depths": [1, 1, 1]},
|
||||
"neck": {"name": "upsample", "channels": 16, "num_upsamples": 2},
|
||||
"head": {"hidden_channels": 16},
|
||||
}
|
||||
|
||||
|
||||
def test_forward_shapes_and_stride(model_cfg):
|
||||
model = build_model(model_cfg).eval()
|
||||
# stem(2) * 2 个下采样 stage = 8，再上采样 4 倍 => stride 2
|
||||
assert model.out_stride == 2
|
||||
|
||||
with torch.no_grad():
|
||||
out = model(torch.randn(2, 3, 64, 64))
|
||||
|
||||
size = 64 // model.out_stride
|
||||
assert out["heatmap"].shape == (2, 3, size, size)
|
||||
assert out["wh"].shape == (2, 2, size, size)
|
||||
assert out["offset"].shape == (2, 2, size, size)
|
||||
assert out["logits"].shape == (2, 3)
|
||||
|
||||
|
||||
def test_backbone_is_convolution_only(model_cfg):
|
||||
model = build_model(model_cfg)
|
||||
forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM, torch.nn.GRU)
|
||||
assert not any(isinstance(m, forbidden) for m in model.modules())
|
||||
|
||||
|
||||
def test_without_classification_head(model_cfg):
|
||||
model = build_model({**model_cfg, "with_classification": False}).eval()
|
||||
with torch.no_grad():
|
||||
out = model(torch.randn(1, 3, 64, 64))
|
||||
assert "logits" not in out
|
||||
|
||||
|
||||
def test_decode_heatmap_recovers_planted_box():
|
||||
heatmap = torch.full((1, 2, 32, 32), -10.0)
|
||||
heatmap[0, 1, 8, 10] = 10.0
|
||||
wh = torch.zeros(1, 2, 32, 32)
|
||||
wh[0, :, 8, 10] = torch.tensor([4.0, 6.0])
|
||||
offset = torch.zeros(1, 2, 32, 32)
|
||||
|
||||
(result,) = decode_heatmap(heatmap, wh, offset, stride=4, topk=5, score_threshold=0.5)
|
||||
|
||||
assert result["labels"].tolist() == [1]
|
||||
assert result["boxes"][0].tolist() == pytest.approx([32.0, 20.0, 48.0, 44.0])
|
||||
|
||||
|
||||
def test_losses_decrease_on_a_single_batch(model_cfg):
|
||||
from engine.losses import DetectionLoss
|
||||
|
||||
torch.manual_seed(0)
|
||||
model = build_model(model_cfg)
|
||||
criterion = DetectionLoss()
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
|
||||
|
||||
images = torch.randn(2, 3, 64, 64)
|
||||
size = 64 // model.out_stride
|
||||
target = {
|
||||
"heatmap": torch.zeros(2, 3, size, size),
|
||||
"wh": torch.zeros(2, 2, size, size),
|
||||
"offset": torch.zeros(2, 2, size, size),
|
||||
"mask": torch.zeros(2, size, size),
|
||||
}
|
||||
target["heatmap"][0, 1, 5, 5] = 1.0
|
||||
target["mask"][0, 5, 5] = 1.0
|
||||
target["wh"][0, :, 5, 5] = 3.0
|
||||
batch = {"target": target, "image_labels": torch.zeros(2, 3)}
|
||||
|
||||
first = last = None
|
||||
for step in range(5):
|
||||
losses = criterion(model(images), batch)
|
||||
optimizer.zero_grad()
|
||||
losses["total"].backward()
|
||||
optimizer.step()
|
||||
total = float(losses["total"].detach())
|
||||
first = total if step == 0 else first
|
||||
last = total
|
||||
|
||||
assert last < first
|
||||
from models.heads.projection import ProjectionHead # noqa: E402
|
||||
|
||||
|
||||
@pytest.mark.parametrize("arch", sorted(_SUPPORTED_ARCHS))
|
||||
@@ -118,10 +34,23 @@ def test_unsupported_arch_is_rejected():
|
||||
BACKBONES.build({"name": "torchvision", "arch": "vit_b_16"})
|
||||
|
||||
|
||||
def test_tiny_conv_shapes():
|
||||
from models.registry import BACKBONES
|
||||
|
||||
backbone = BACKBONES.build(
|
||||
{"name": "tiny_conv", "widths": [8, 16, 32], "depths": [1, 1, 1]}
|
||||
).eval()
|
||||
with torch.no_grad():
|
||||
feat = backbone(torch.randn(2, 3, 64, 64))
|
||||
assert feat.shape[:2] == (2, 32)
|
||||
assert backbone.out_channels == 32
|
||||
# stem 下采样 2 倍，其后每个 stage(除第一个)各下采样 2 倍
|
||||
assert backbone.out_stride == 8
|
||||
|
||||
|
||||
@pytest.mark.parametrize("arch", ["resnet18", "shufflenet_v2_x1_0", "mobilenet_v3_small"])
|
||||
def test_alternative_backbones_drive_simclr(arch):
|
||||
"""换骨架后 SimCLR 的投影头要能接上（out_channels 必须准）。"""
|
||||
from models.heads.projection import ProjectionHead
|
||||
from models.registry import BACKBONES
|
||||
|
||||
backbone = BACKBONES.build({"name": "torchvision", "arch": arch, "pretrained": False}).eval()
|
||||
@@ -130,18 +59,13 @@ def test_alternative_backbones_drive_simclr(arch):
|
||||
assert head(backbone(torch.randn(2, 3, 64, 64))).shape == (2, 16)
|
||||
|
||||
|
||||
def test_alternative_backbone_drives_defect_net():
|
||||
"""检测模型换骨架同样要能跑通（neck 的 in_channels 由 backbone 决定）。"""
|
||||
model = build_model(
|
||||
{
|
||||
"name": "defect_net",
|
||||
"num_classes": 3,
|
||||
"backbone": {"name": "torchvision", "arch": "shufflenet_v2_x1_0", "pretrained": False},
|
||||
"neck": {"name": "upsample", "channels": 16, "num_upsamples": 3},
|
||||
"head": {"hidden_channels": 16},
|
||||
}
|
||||
).eval()
|
||||
assert model.out_stride == 4
|
||||
def test_classifier_shapes_and_convolution_only():
|
||||
model = build_model({
|
||||
"name": "classifier", "num_classes": 4,
|
||||
"backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
|
||||
"classification": {"dropout": 0.1},
|
||||
}).eval()
|
||||
with torch.no_grad():
|
||||
out = model(torch.randn(1, 3, 64, 64))
|
||||
assert out["heatmap"].shape == (1, 3, 16, 16)
|
||||
assert model(torch.randn(3, 3, 64, 64)).shape == (3, 4)
|
||||
forbidden = (torch.nn.MultiheadAttention, torch.nn.Transformer, torch.nn.LSTM)
|
||||
assert not any(isinstance(m, forbidden) for m in model.modules())
|
||||
|
||||
Reference in New Issue
Block a user