| """One-off: InsightFace's ONNX files -> safetensors for nets.py. Needs onnx (only here). |
| |
| python convert.py <models dir> <out dir> |
| """ |
|
|
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| import torch |
| from onnx import helper, numpy_helper |
| from safetensors.torch import save_file |
| from torch import nn |
|
|
| from nets import SCRFD, IResNet |
|
|
| PARAM_OPS = {"Conv": nn.Conv2d, "BatchNormalization": (nn.BatchNorm2d, nn.BatchNorm1d), "PRelu": nn.PReLU, "Gemm": nn.Linear} |
|
|
|
|
| def attrs(n): |
| return {a.name: helper.get_attribute_value(a) for a in n.attribute} |
|
|
|
|
| def load(model: nn.Module, path: Path) -> nn.Module: |
| g = onnx.load(str(path)).graph |
| init = {t.name: torch.from_numpy(numpy_helper.to_array(t).copy()) for t in g.initializer} |
| mods = [m for m in model.modules() if isinstance(m, (nn.Conv2d, nn.BatchNorm2d, nn.BatchNorm1d, nn.PReLU, nn.Linear))] |
| scales = [m for m in model.modules() if hasattr(m, "scale") and isinstance(getattr(m, "scale"), torch.Tensor)] |
| nodes = [n for n in g.node if n.op_type in PARAM_OPS] |
| assert len(nodes) == len(mods), (len(nodes), len(mods)) |
| for n, m in zip(nodes, mods): |
| assert isinstance(m, PARAM_OPS[n.op_type]), (n.name, n.op_type, m) |
| w = [init[i] for i in n.input[1:]] |
| a = attrs(n) |
| if n.op_type == "Conv": |
| assert a.get("group", 1) == 1 and a.get("dilations", [1, 1]) == [1, 1] |
| assert list(a["strides"]) == list(m.stride) and list(a["pads"]) == list(m.padding) * 2, (n.name, a, m) |
| m.weight.data.copy_(w[0]), m.bias.data.copy_(w[1]) |
| elif n.op_type == "BatchNormalization": |
| assert abs(a["epsilon"] - m.eps) < 1e-9 |
| for t, v in zip((m.weight, m.bias, m.running_mean, m.running_var), w): |
| t.data.copy_(v) |
| elif n.op_type == "PRelu": |
| m.weight.data.copy_(w[0].reshape(-1)) |
| else: |
| assert a.get("transB") == 1 and a.get("alpha", 1.0) == 1.0 and a.get("beta", 1.0) == 1.0 |
| m.weight.data.copy_(w[0]), m.bias.data.copy_(w[1]) |
| muls = [n for n in g.node if n.op_type == "Mul"] |
| assert len(muls) == len(scales) |
| for n, s in zip(muls, scales): |
| s.scale.copy_(init[n.input[1]]) |
| |
| for n in g.node: |
| if n.op_type == "Resize": |
| print(" resize", {k: (v.decode() if isinstance(v, bytes) else v) for k, v in attrs(n).items()}) |
| if n.op_type == "Reshape": |
| assert init[n.input[1]].tolist() in ([-1, 1], [-1, 4], [-1, 10]), init[n.input[1]] |
| if muls: |
| print(" head scales", [float(s.scale) for s in scales]) |
| return model.eval() |
|
|
|
|
| def main(src: Path, out: Path): |
| out.mkdir(parents=True, exist_ok=True) |
| for name, model, file in [("det_10g", SCRFD(), "buffalo_l/det_10g.onnx"), |
| ("w600k_r50", IResNet(IResNet.LAYERS["r50"]), "buffalo_l/w600k_r50.onnx"), |
| ("glintr100", IResNet(IResNet.LAYERS["r100"]), "antelopev2/antelopev2/glintr100.onnx")]: |
| print(name) |
| load(model, src / file) |
| |
| |
| dtype = torch.float32 if isinstance(model, SCRFD) else torch.float16 |
| save_file({k: v.to(dtype).contiguous() for k, v in model.state_dict().items() if not k.endswith("num_batches_tracked")}, |
| str(out / f"{name}.safetensors"), metadata={"source": file, "converted_by": "ZenTrainer convert.py"}) |
|
|
|
|
| if __name__ == "__main__": |
| main(Path(sys.argv[1]), Path(sys.argv[2])) |
|
|
