Add export to onnx and tensorRT

scripts/export_onnx_trt.py

@@ -0,0 +1,130 @@
+import argparse
+import os
+import pprint
+from typing import Tuple
+
+import torch
+
+from src.wireseghr.model import WireSegHR
+
+
+class CoarseModule(torch.nn.Module):
+    def __init__(self, core: WireSegHR):
+        super().__init__()
+        self.core = core
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        logits, cond = self.core.forward_coarse(x)
+        return logits, cond
+
+
+class FineModule(torch.nn.Module):
+    def __init__(self, core: WireSegHR):
+        super().__init__()
+        self.core = core
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        logits = self.core.forward_fine(x)
+        return logits
+
+
+def build_model(cfg: dict, device: torch.device) -> WireSegHR:
+    pretrained_flag = bool(cfg.get("pretrained", False))
+    model = WireSegHR(backbone=cfg["backbone"], in_channels=6, pretrained=pretrained_flag)
+    model = model.to(device)
+    return model
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Export WireSegHR to ONNX and TensorRT")
+    parser.add_argument("--config", type=str, default="configs/default.yaml")
+    parser.add_argument("--ckpt", type=str, default="", help="Path to checkpoint .pt")
+    parser.add_argument("--out_dir", type=str, default="exports")
+    parser.add_argument("--coarse_size", type=int, default=1024)
+    parser.add_argument("--fine_patch_size", type=int, default=1024)
+    parser.add_argument("--opset", type=int, default=17)
+    parser.add_argument("--trtexec", type=str, default="", help="Optional path to trtexec to build TRT engines")
+
+    args = parser.parse_args()
+
+    import yaml
+
+    with open(args.config, "r") as f:
+        cfg = yaml.safe_load(f)
+    print("[export] Loaded config:")
+    pprint.pprint(cfg)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = build_model(cfg, device)
+
+    ckpt_path = args.ckpt if args.ckpt else cfg.get("resume", "")
+    if ckpt_path:
+        assert os.path.isfile(ckpt_path), f"Checkpoint not found: {ckpt_path}"
+        print(f"[export] Loading checkpoint: {ckpt_path}")
+        state = torch.load(ckpt_path, map_location=device)
+        model.load_state_dict(state["model"])  # expects dict with key 'model'
+    model.eval()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    # Prepare dummy inputs (static shapes for best TRT performance)
+    coarse_in = torch.randn(1, 6, args.coarse_size, args.coarse_size, device=device)
+    fine_in = torch.randn(1, 6, args.fine_patch_size, args.fine_patch_size, device=device)
+
+    # Coarse export
+    coarse_wrapper = CoarseModule(model).to(device).eval()
+    coarse_onnx = os.path.join(args.out_dir, f"wireseghr_coarse_{args.coarse_size}.onnx")
+    print(f"[export] Exporting COARSE to {coarse_onnx}")
+    torch.onnx.export(
+        coarse_wrapper,
+        coarse_in,
+        coarse_onnx,
+        export_params=True,
+        opset_version=args.opset,
+        do_constant_folding=True,
+        input_names=["x_coarse"],
+        output_names=["logits", "cond"],
+        dynamic_axes=None,
+    )
+
+    # Fine export
+    fine_wrapper = FineModule(model).to(device).eval()
+    fine_onnx = os.path.join(args.out_dir, f"wireseghr_fine_{args.fine_patch_size}.onnx")
+    print(f"[export] Exporting FINE to {fine_onnx}")
+    torch.onnx.export(
+        fine_wrapper,
+        fine_in,
+        fine_onnx,
+        export_params=True,
+        opset_version=args.opset,
+        do_constant_folding=True,
+        input_names=["x_fine"],
+        output_names=["logits"],
+        dynamic_axes=None,
+    )
+
+    # Optional TensorRT building via trtexec
+    if args.trtexec:
+        import subprocess
+
+        def build_engine(onnx_path: str, engine_path: str):
+            print(f"[export] Building TRT engine: {engine_path}")
+            cmd = [
+                args.trtexec,
+                f"--onnx={onnx_path}",
+                f"--saveEngine={engine_path}",
+                "--explicitBatch",
+                "--fp16",
+            ]
+            subprocess.run(cmd, check=True)
+
+        coarse_engine = os.path.join(args.out_dir, f"wireseghr_coarse_{args.coarse_size}.engine")
+        fine_engine = os.path.join(args.out_dir, f"wireseghr_fine_{args.fine_patch_size}.engine")
+        build_engine(coarse_onnx, coarse_engine)
+        build_engine(fine_onnx, fine_engine)
+
+    print("[export] Done.")
+
+
+if __name__ == "__main__":
+    main()