feat: train fixed-size classifiers and export class probabilities
This commit is contained in:
Shaun Zhou
@@ -0,0 +1,134 @@
|
||||
#!/usr/bin/env python
|
||||
"""Compare supervised classification accuracy, export parity and YOLO agreement.
|
||||
|
||||
Uses fixed-size human-labelled validation tiles, never training data. No resize.
|
||||
Empty detector output is normal; the strongest valid candidate supplies top1.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from datasets.tile_classification import YoloTileClassificationDataset # noqa: E402
|
||||
from models import build_model # noqa: E402
|
||||
|
||||
|
||||
def output_class(output: np.ndarray, conf: float, background: int) -> int:
|
||||
scores = output[0, 4:, :]
|
||||
index = int(scores.max(axis=0).argmax())
|
||||
return int(scores[:, index].argmax()) if scores[:, index].max() >= conf else background
|
||||
|
||||
|
||||
def probability_class(output: np.ndarray, conf: float, background: int) -> int:
|
||||
if output.ndim != 2 or output.shape[0] != 1:
|
||||
raise ValueError("Expected classifier probabilities [1,C]")
|
||||
return int(output[0].argmax()) if output[0].max() >= conf else background
|
||||
|
||||
|
||||
def matrix_report(matrix, names, background):
|
||||
counts = matrix.sum(1)
|
||||
foreground = int(matrix.sum() - counts[background])
|
||||
return {"agreement": float(np.trace(matrix) / matrix.sum()),
|
||||
"foreground_agreement": float((np.trace(matrix) - matrix[background, background])
|
||||
/ foreground) if foreground else None,
|
||||
"per_reference_class": {name: {"support": int(counts[i]),
|
||||
"agreement": float(matrix[i, i] / counts[i])
|
||||
if counts[i] else None}
|
||||
for i, name in enumerate(names)},
|
||||
"confusion_reference_rows_prediction_columns": matrix.tolist()}
|
||||
|
||||
|
||||
def main():
|
||||
import openvino as ov
|
||||
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--teacher", required=True, help="Reference YOLO .xml; used only for comparison")
|
||||
parser.add_argument("--student", required=True, help="Classifier probability .xml, output [1,C]")
|
||||
parser.add_argument("--checkpoint", required=True)
|
||||
parser.add_argument("--report", required=True)
|
||||
parser.add_argument("--conf", type=float, default=0.25)
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
args = parser.parse_args()
|
||||
if not 0 < args.conf < 1:
|
||||
raise ValueError("Positive confidence threshold required")
|
||||
state = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
|
||||
cfg = state["config"]
|
||||
names = cfg["data"]["classes"]
|
||||
background = names.index("normal")
|
||||
for xml in (args.teacher, args.student):
|
||||
meta = yaml.safe_load((Path(xml).parent / "metadata.yaml").read_text())
|
||||
if [meta["names"][i] for i in range(len(meta["names"]))] != names:
|
||||
raise ValueError(f"Class order mismatch: {xml}")
|
||||
if cfg["data"]["val_transforms"] != [{"name": "ToTensor"}]:
|
||||
raise ValueError("This comparison expects the new raw RGB / 255 tile classifiers")
|
||||
ds = YoloTileClassificationDataset(cfg["data"]["yaml"], "val", cfg["data"]["image_size"],
|
||||
transforms=[{"name": "ToTensor"}], classes=names)
|
||||
model_cfg = copy.deepcopy(cfg["model"])
|
||||
model_cfg["backbone"]["pretrained"] = False
|
||||
model_cfg["init_from"] = None
|
||||
model_cfg["num_classes"] = len(names)
|
||||
model = build_model(model_cfg).to(args.device).eval()
|
||||
model.load_state_dict(state["model"])
|
||||
core = ov.Core()
|
||||
settings = {"PERFORMANCE_HINT": "LATENCY", "INFERENCE_NUM_THREADS": 4}
|
||||
teacher = core.compile_model(args.teacher, "CPU", settings)
|
||||
student = core.compile_model(args.student, "CPU", settings)
|
||||
for compiled in (teacher, student):
|
||||
if list(compiled.input().shape) != [1, 3, *cfg["data"]["image_size"]]:
|
||||
raise ValueError("Unexpected model interface")
|
||||
if (list(teacher.output().shape)[:2] != [1, 4 + len(names)]
|
||||
or list(student.output().shape) != [1, len(names)]):
|
||||
raise ValueError("Expected raw YOLO output and plain classifier probabilities")
|
||||
matrices = {key: np.zeros((len(names), len(names)), dtype=int)
|
||||
for key in ["pt_accuracy", "ov_accuracy", "ov_deployment_accuracy",
|
||||
"yolo_deployment_accuracy",
|
||||
"export_top1_parity", "export_deployment_parity", "yolo_student_agreement"]}
|
||||
records = []
|
||||
with torch.inference_mode():
|
||||
for start in range(0, len(ds), 64):
|
||||
batch = [ds[i] for i in range(start, min(start + 64, len(ds)))]
|
||||
images = torch.stack([b["image"] for b in batch])
|
||||
probabilities = model(images.to(args.device)).softmax(1).cpu().numpy()
|
||||
for i, item in enumerate(batch):
|
||||
x = images[i:i+1].numpy()
|
||||
t = output_class(teacher(x)[teacher.output()], args.conf, background)
|
||||
raw = student(x)[student.output()]
|
||||
if (not np.isfinite(raw).all() or (raw < 0).any() or (raw > 1).any()
|
||||
or not np.allclose(raw.sum(1), 1, atol=1e-5)):
|
||||
raise ValueError("Invalid classifier probability output")
|
||||
s = probability_class(raw, args.conf, background)
|
||||
raw_s = int(raw[0].argmax())
|
||||
p = int(probabilities[i].argmax())
|
||||
deployed_p = p if probabilities[i].max() >= args.conf else background
|
||||
y = item["label"]
|
||||
for key, row, col in [("pt_accuracy", y, p), ("ov_accuracy", y, raw_s),
|
||||
("ov_deployment_accuracy", y, s),
|
||||
("yolo_deployment_accuracy", y, t),
|
||||
("export_top1_parity", p, raw_s),
|
||||
("export_deployment_parity", deployed_p, s),
|
||||
("yolo_student_agreement", t, s)]:
|
||||
matrices[key][row, col] += 1
|
||||
records.append({"path": item["meta"]["path"], "truth": y, "pytorch": p,
|
||||
"pytorch_deployed": deployed_p, "openvino": s,
|
||||
"openvino_top1": raw_s, "yolo": t,
|
||||
"student_max_score": float(raw.max())})
|
||||
report = {"classes": names, "images": len(ds), "conf": args.conf,
|
||||
"teacher": args.teacher, "student": args.student, "checkpoint": args.checkpoint,
|
||||
"metrics": {k: matrix_report(m, names, background) for k, m in matrices.items()}}
|
||||
out = Path(args.report)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
out.write_text(json.dumps(report, indent=2))
|
||||
out.with_suffix(".predictions.json").write_text(json.dumps(records, indent=2))
|
||||
print(json.dumps(report, indent=2), flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,10 +1,8 @@
|
||||
#!/usr/bin/env python
|
||||
"""Export a classifier with the input/output contract of a YOLO OpenVINO IR.
|
||||
"""Export RGB / 255 BCHW -> Softmax probabilities [1, C], in training class order.
|
||||
|
||||
The sole active candidate is a full-image pseudo box, not a localization result.
|
||||
Input must already be RGB float32 BCHW in [0, 1], exactly as in deploy/cpp.
|
||||
Training normalization is embedded. No resize, pseudo boxes or background gating.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
@@ -19,7 +17,6 @@ import yaml
|
||||
from torch import nn
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
from models import build_model # noqa: E402
|
||||
|
||||
|
||||
@@ -36,63 +33,48 @@ def ordered_names(names) -> list[str]:
|
||||
return names
|
||||
|
||||
|
||||
def class_order(source: list[str], target: list[str]) -> list[int]:
|
||||
source, target = ordered_names(source), ordered_names(target)
|
||||
if set(source) != set(target):
|
||||
raise ValueError(f"Class sets differ: classifier={source}, YOLO={target}")
|
||||
return [source.index(name) for name in target]
|
||||
def input_size(config: dict) -> list[int]:
|
||||
size = config["data"]["image_size"]
|
||||
size = [size, size] if isinstance(size, int) else list(size)
|
||||
if len(size) != 2 or any(not isinstance(v, int) or v < 1 for v in size):
|
||||
raise ValueError(f"Invalid training image_size: {size}")
|
||||
return size
|
||||
|
||||
|
||||
def preprocessing(config: dict, size: int) -> float | None:
|
||||
"""Reject preprocessing that cannot be reproduced without resampling the input."""
|
||||
def preprocessing(config: dict, size: int | list[int]) -> float | None:
|
||||
shape = [size, size] if isinstance(size, int) else list(size)
|
||||
transforms = config["data"]["val_transforms"]
|
||||
sequence = [t["name"] for t in transforms]
|
||||
if sequence not in (["Resize", "ToTensor", "PerImageStandardize"],
|
||||
["ToTensor", "PerImageStandardize"], ["Resize", "ToTensor"], ["ToTensor"]):
|
||||
raise ValueError(f"Unsupported validation preprocessing: {sequence}")
|
||||
eps = None
|
||||
for transform in transforms:
|
||||
if transform["name"] == "Resize":
|
||||
if transform["size"] != [size, size]:
|
||||
raise ValueError("Training size differs from template; retrain at the target size")
|
||||
if transform["name"] == "Resize" and list(transform["size"]) != shape:
|
||||
raise ValueError("Training size differs from export input")
|
||||
if transform["name"] == "PerImageStandardize":
|
||||
return float(transform.get("eps", 1e-5))
|
||||
return None
|
||||
eps = float(transform.get("eps", 1e-5))
|
||||
if eps <= 0:
|
||||
raise ValueError("Standardization eps must be positive")
|
||||
return eps
|
||||
|
||||
|
||||
class YoloClassifierAdapter(nn.Module):
|
||||
def __init__(self, classifier: nn.Module, size: int, anchors: int,
|
||||
order: list[int], eps: float | None = 1e-5,
|
||||
background_index: int | None = None):
|
||||
class ClassifierProbabilities(nn.Module):
|
||||
def __init__(self, classifier: nn.Module, eps: float | None = None):
|
||||
super().__init__()
|
||||
if size <= 0 or anchors <= 0:
|
||||
raise ValueError("size and anchors must be positive")
|
||||
self.classifier = classifier
|
||||
self.eps = eps
|
||||
if background_index is not None and not 0 <= background_index < len(order):
|
||||
raise ValueError("background_index must be an output class index")
|
||||
self.background_index = background_index
|
||||
self.register_buffer("order", torch.tensor(order, dtype=torch.long))
|
||||
self.register_buffer("box", torch.tensor([size / 2, size / 2, size, size])
|
||||
.reshape(1, 4, 1))
|
||||
# Only candidate zero is active; do not duplicate predictions across anchors.
|
||||
self.register_buffer("inactive", torch.zeros(1, 4 + len(order), anchors - 1))
|
||||
self.classifier, self.eps = classifier, eps
|
||||
|
||||
def normalize(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self.eps is None:
|
||||
return x
|
||||
dims = (1, 2, 3)
|
||||
mean = x.mean(dim=dims, keepdim=True)
|
||||
# Match datasets.transforms.PerImageStandardize: torch.std correction=1.
|
||||
# Match training torch.std(correction=1), separately for each image.
|
||||
variance = ((x - mean) ** 2).sum(dim=dims, keepdim=True) / (x[0].numel() - 1)
|
||||
return (x - mean) / (torch.sqrt(variance) + self.eps)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
scores = self.classifier(self.normalize(x)).softmax(dim=1)[:, self.order]
|
||||
if self.background_index is not None:
|
||||
foreground = (scores.argmax(dim=1) != self.background_index).unsqueeze(1)
|
||||
scores = scores * foreground.to(scores.dtype)
|
||||
active = torch.cat((self.box.expand(x.shape[0], -1, -1), scores.unsqueeze(2)), dim=1)
|
||||
return torch.cat((active, self.inactive.expand(x.shape[0], -1, -1)), dim=2)
|
||||
return self.classifier(self.normalize(x)).softmax(dim=1)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
@@ -100,97 +82,71 @@ def main() -> int:
|
||||
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--checkpoint", default="runs/classify/best.pth")
|
||||
parser.add_argument("--template", required=True, help="Production YOLO .xml (metadata.yaml beside it)")
|
||||
parser.add_argument("--classes", help="Classifier class order: classes.txt or eval_report.json")
|
||||
parser.add_argument("--background-class", default="normal", help="Top1 class to emit no boxes for")
|
||||
parser.add_argument("--keep-background-box", action="store_true",
|
||||
help="Opt out of background suppression (legacy classification behavior)")
|
||||
parser.add_argument("--classes", help="Training class order: classes.txt or eval_report.json")
|
||||
parser.add_argument("--fp32-weights", action="store_true", help="Keep full-precision weights")
|
||||
parser.add_argument("--out", required=True, help="New directory; existing directories are rejected")
|
||||
args = parser.parse_args()
|
||||
checkpoint, template, out = Path(args.checkpoint), Path(args.template), Path(args.out)
|
||||
checkpoint, out = Path(args.checkpoint), Path(args.out)
|
||||
if out.exists():
|
||||
raise FileExistsError(f"Refusing to overwrite {out}")
|
||||
state = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
||||
config = state["config"]
|
||||
names = config["data"].get("classes")
|
||||
if args.classes:
|
||||
path = Path(args.classes)
|
||||
source = (json.loads(path.read_text())["classes"] if path.suffix == ".json"
|
||||
else path.read_text().splitlines())
|
||||
else:
|
||||
source = config["data"].get("classes")
|
||||
if not source:
|
||||
report = checkpoint.parent / "eval_report.json"
|
||||
if not report.exists():
|
||||
raise ValueError("No recorded class order; provide --classes (do not guess folder order)")
|
||||
source = json.loads(report.read_text())["classes"]
|
||||
source = ordered_names(source)
|
||||
metadata = yaml.safe_load((template.parent / "metadata.yaml").read_text())
|
||||
if metadata.get("end2end") or metadata.get("args", {}).get("nms"):
|
||||
raise ValueError("Expected raw [1, 4+nc, anchors] output, not end-to-end/NMS output")
|
||||
target = ordered_names(metadata["names"])
|
||||
order = class_order(source, target)
|
||||
if not args.keep_background_box and args.background_class not in target:
|
||||
raise ValueError(f"Background class {args.background_class!r} is not in {target}")
|
||||
background_index = None if args.keep_background_box else target.index(args.background_class)
|
||||
core = ov.Core()
|
||||
reference = core.read_model(str(template))
|
||||
if len(reference.inputs) != 1 or len(reference.outputs) != 1:
|
||||
raise ValueError("Expected one input and one output")
|
||||
input_shape, output_shape = list(reference.input().shape), list(reference.output().shape)
|
||||
if (len(input_shape) != 4 or input_shape[:2] != [1, 3]
|
||||
or input_shape[2] != input_shape[3] or len(output_shape) != 3
|
||||
or output_shape[:2] != [1, 4 + len(target)] or output_shape[2] < 1):
|
||||
raise ValueError(f"Unsupported template shapes: {input_shape}, {output_shape}")
|
||||
if (reference.input().element_type != ov.Type.f32
|
||||
or reference.output().element_type != ov.Type.f32):
|
||||
raise ValueError("Only float32 input/output are supported")
|
||||
size = input_shape[2]
|
||||
supplied = (json.loads(path.read_text())["classes"] if path.suffix == ".json"
|
||||
else path.read_text().splitlines())
|
||||
if names and ordered_names(supplied) != ordered_names(names):
|
||||
raise ValueError("--classes contradicts checkpoint class order")
|
||||
names = supplied
|
||||
if not names:
|
||||
report = checkpoint.parent / "eval_report.json"
|
||||
if not report.exists():
|
||||
raise ValueError("No recorded class order; provide --classes (do not guess folder order)")
|
||||
names = json.loads(report.read_text())["classes"]
|
||||
names = ordered_names(names)
|
||||
size = input_size(config)
|
||||
eps = preprocessing(config, size)
|
||||
model_config = copy.deepcopy(config["model"])
|
||||
model_config["backbone"]["pretrained"] = False
|
||||
model_config["init_from"] = None
|
||||
model_config["num_classes"] = len(source)
|
||||
model_config["num_classes"] = len(names)
|
||||
model = build_model(model_config).eval()
|
||||
model.load_state_dict(state["model"], strict=True)
|
||||
adapter = YoloClassifierAdapter(model, size, output_shape[2], order, eps,
|
||||
background_index=background_index).eval()
|
||||
wrapper = ClassifierProbabilities(model, eps).eval()
|
||||
shape = [1, model_config["backbone"].get("in_channels", 3), *size]
|
||||
torch.manual_seed(42)
|
||||
example = torch.rand(input_shape)
|
||||
converted = ov.convert_model(adapter, example_input=example, input=input_shape)
|
||||
converted.input().get_tensor().set_names(reference.input().get_names())
|
||||
converted.output().get_tensor().set_names(reference.output().get_names())
|
||||
# Preserve any interface metadata used by OpenVINO consumers, then mark pseudo boxes.
|
||||
for key, value in reference.get_rt_info().items():
|
||||
if key != "conversion_parameters":
|
||||
converted.set_rt_info(value, [key])
|
||||
converted.set_rt_info("classifier_pseudo_box", ["adapter_kind"])
|
||||
example = torch.rand(shape)
|
||||
converted = ov.convert_model(wrapper, example_input=example, input=shape)
|
||||
converted.input().get_tensor().set_names({"images"})
|
||||
converted.output().get_tensor().set_names({"probabilities"})
|
||||
converted.set_rt_info("classify", ["model_info", "task"])
|
||||
converted.set_rt_info(json.dumps(dict(enumerate(names))), ["model_info", "names"])
|
||||
out.mkdir(parents=True)
|
||||
xml = out / template.name
|
||||
ov.save_model(converted, str(xml), compress_to_fp16=True)
|
||||
compiled = core.compile_model(str(xml), "CPU", {"PERFORMANCE_HINT": "LATENCY"})
|
||||
arch = model_config["backbone"].get("arch", model_config["backbone"]["name"])
|
||||
xml = out / f"{arch}.xml"
|
||||
ov.save_model(converted, str(xml), compress_to_fp16=not args.fp32_weights)
|
||||
compiled = ov.Core().compile_model(str(xml), "CPU", {"PERFORMANCE_HINT": "LATENCY"})
|
||||
errors = []
|
||||
with torch.inference_mode():
|
||||
for sample in (example, torch.zeros_like(example), torch.ones_like(example)):
|
||||
expected = adapter(sample).numpy()
|
||||
expected = wrapper(sample).numpy()
|
||||
actual = compiled(sample.numpy())[compiled.output()]
|
||||
assert actual.shape == (1, len(names))
|
||||
assert np.isfinite(actual).all() and (actual >= 0).all() and (actual <= 1).all()
|
||||
np.testing.assert_allclose(actual.sum(axis=1), 1, atol=1e-5)
|
||||
np.testing.assert_allclose(actual, expected, atol=0.02, rtol=0.02)
|
||||
assert np.argmax(actual[0, 4:, 0]) == np.argmax(expected[0, 4:, 0])
|
||||
assert np.count_nonzero(actual[:, :, 1:]) == 0
|
||||
assert actual.argmax(axis=1).tolist() == expected.argmax(axis=1).tolist()
|
||||
errors.append(float(np.max(np.abs(expected - actual))))
|
||||
metadata.update({"names": dict(enumerate(target)), "adapter_kind": "classifier_pseudo_box",
|
||||
"description": "Classifier with one full-image pseudo box; no localization",
|
||||
"background_class": None if background_index is None else args.background_class,
|
||||
"end2end": False, "args": {"nms": False}, "imgsz": [size, size]})
|
||||
metadata = {"task": "classify", "names": dict(enumerate(names)), "imgsz": size,
|
||||
"channels": shape[1], "layout": "BCHW", "input_range": [0, 1],
|
||||
"output": "softmax_probabilities", "input_shape": shape,
|
||||
"output_shape": [1, len(names)], "standardize_eps": eps,
|
||||
"weight_storage": "fp32" if args.fp32_weights else "fp16"}
|
||||
(out / "metadata.yaml").write_text(yaml.safe_dump(metadata, allow_unicode=True))
|
||||
(out / "classes.txt").write_text("\n".join(target) + "\n")
|
||||
summary = {"checkpoint": str(checkpoint.resolve()), "template": str(template.resolve()),
|
||||
"classifier_classes": source, "output_classes": target, "class_order": order,
|
||||
"input_shape": input_shape, "output_shape": output_shape,
|
||||
"standardize_eps": eps, "max_abs_errors": errors,
|
||||
"background_class": None if background_index is None else args.background_class,
|
||||
"background_index": background_index,
|
||||
"pseudo_box": [size / 2, size / 2, size, size]}
|
||||
(out / "adapter.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
(out / "classes.txt").write_text("\n".join(names) + "\n")
|
||||
summary = {"checkpoint": str(checkpoint.resolve()), **metadata, "max_abs_errors": errors}
|
||||
(out / "export.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
print(json.dumps({"model": str(xml), **summary}, ensure_ascii=False, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
Reference in New Issue
Block a user