156 lines
7.1 KiB
Python
156 lines
7.1 KiB
Python
#!/usr/bin/env python
|
|
"""Export RGB / 255 BCHW -> Softmax probabilities [1, C], in training class order.
|
|
|
|
Training normalization is embedded. No resize, pseudo boxes or background gating.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import yaml
|
|
from torch import nn
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
from models import build_model # noqa: E402
|
|
|
|
|
|
def ordered_names(names) -> list[str]:
|
|
if isinstance(names, dict):
|
|
indexed = {int(k): v for k, v in names.items()}
|
|
if set(indexed) != set(range(len(indexed))):
|
|
raise ValueError("Class IDs must be contiguous from zero")
|
|
names = [indexed[i] for i in range(len(indexed))]
|
|
if not isinstance(names, list) or not names or not all(isinstance(n, str) for n in names):
|
|
raise ValueError("Expected a nonempty class-name list or indexed mapping")
|
|
if len(set(names)) != len(names):
|
|
raise ValueError("Duplicate class names")
|
|
return names
|
|
|
|
|
|
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 | 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" and list(transform["size"]) != shape:
|
|
raise ValueError("Training size differs from export input")
|
|
if transform["name"] == "PerImageStandardize":
|
|
eps = float(transform.get("eps", 1e-5))
|
|
if eps <= 0:
|
|
raise ValueError("Standardization eps must be positive")
|
|
return eps
|
|
|
|
|
|
class ClassifierProbabilities(nn.Module):
|
|
def __init__(self, classifier: nn.Module, eps: float | None = None):
|
|
super().__init__()
|
|
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 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:
|
|
return self.classifier(self.normalize(x)).softmax(dim=1)
|
|
|
|
|
|
def main() -> int:
|
|
import openvino as ov
|
|
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--checkpoint", default="runs/classify/best.pth")
|
|
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, 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)
|
|
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(names)
|
|
model = build_model(model_config).eval()
|
|
model.load_state_dict(state["model"], strict=True)
|
|
wrapper = ClassifierProbabilities(model, eps).eval()
|
|
shape = [1, model_config["backbone"].get("in_channels", 3), *size]
|
|
torch.manual_seed(42)
|
|
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)
|
|
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 = 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 actual.argmax(axis=1).tolist() == expected.argmax(axis=1).tolist()
|
|
errors.append(float(np.max(np.abs(expected - actual))))
|
|
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(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
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|