""" ONNX export scripts for WeDetect. Two separate exports: 1. Image Encoder + Neck + Head → wedetect_image_head.onnx 2. Text Encoder → wedetect_text_encoder.onnx Usage: python export_onnx.py --config config/wedetect_base.py --checkpoint checkpoints/wedetect_base.pth The text encoder is exported separately so text embeddings can be pre-computed offline for any category set. At inference time the image+head ONNX model takes the pre-computed text features together with the image and produces detections. """ import argparse import os.path as osp import warnings from io import BytesIO import onnx import torch import torch.nn as nn from mmdet.utils import register_all_modules register_all_modules() from mmengine.config import Config from mmdet.apis import init_detector warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) warnings.filterwarnings(action='ignore', category=torch.jit.ScriptWarning) warnings.filterwarnings(action='ignore', category=UserWarning) warnings.filterwarnings(action='ignore', category=FutureWarning) # --------------------------------------------------------------------------- # Wrapper: Image Encoder + Neck + Head # --------------------------------------------------------------------------- class ImageHeadWrapper(nn.Module): """Pack ConvNeXt → Neck → Head into a single traceable module. Inputs: image [B, 3, 640, 640] float32 text_features [B, num_classes, 768] float32 Outputs (three scales, strides 8 / 16 / 32): cls_scores_s8 [B, num_classes, 80, 80] cls_scores_s16 [B, num_classes, 40, 40] cls_scores_s32 [B, num_classes, 20, 20] bbox_preds_s8 [B, 4, 80, 80] bbox_preds_s16 [B, 4, 40, 40] bbox_preds_s32 [B, 4, 20, 20] """ def __init__(self, model): super().__init__() self.backbone = model.backbone self.neck = model.neck self.head_module = model.bbox_head.head_module def forward(self, image, text_features): # ConvNeXt (forward_image skips the text encoder) img_feats = self.backbone.forward_image(image) # BiFPN neck if self.neck is not None: img_feats = self.neck(img_feats) # Head: returns ((cls_L0, cls_L1, cls_L2), (bbox_L0, bbox_L1, bbox_L2)) cls_scores, bbox_preds = self.head_module(img_feats, text_features) # Flatten to a single tuple so ONNX output names align correctly return ( cls_scores[0], bbox_preds[0], # stride 8 cls_scores[1], bbox_preds[1], # stride 16 cls_scores[2], bbox_preds[2], # stride 32 ) # --------------------------------------------------------------------------- # Wrapper: Text Encoder (XLM-RoBERTa without tokenizer) # --------------------------------------------------------------------------- class TextEncoderWrapper(nn.Module): """Wrap the XLM-RoBERTa backbone so it can be exported without a tokenizer. Inputs: input_ids [num_texts, max_seq_len] int64 (fixed seq_len) attention_mask [num_texts, max_seq_len] int64 Output: text_features [1, num_texts, 768] float32 (L2-normalised) """ def __init__(self, text_model): super().__init__() self.model = text_model.model # XLMRobertaModel self.head = text_model.head # Linear(768 → 768) def forward(self, input_ids, attention_mask): # Compute position_ids in float32 to avoid INT32 CumSum/Mul/Add ops # that the NPU backend does not support. # XLMRoberta uses padding_idx=1, so mask = (input_ids != 1). mask = input_ids.ne(1).float() # [N, L] float32 position_ids = torch.cumsum(mask, dim=1) * mask # cumsum+mul in float32 position_ids = (position_ids + 1.0).long() # +padding_idx in float, then cast out = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, ) cls_embed = out["last_hidden_state"][:, 0, :] # [N, 768] txt_feats = self.head(cls_embed) # L2-normalise + add batch dim → (1, N, 768) txt_feats = nn.functional.normalize(txt_feats, dim=-1) return txt_feats.unsqueeze(0) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _onnx_simplify(onnx_model, output_path): """Save ONNX, then try to simplify and overwrite. Saving first guarantees a usable file even if onnxsim crashes. """ # Save unsimplified first — always produces a valid file on disk onnx.save(onnx_model, output_path) print(f"Saved (unsimplified): {output_path}") # Then attempt simplify try: import onnxsim simplified, check = onnxsim.simplify(onnx_model) if check: onnx.save(simplified, output_path) print("ONNX simplify: passed, overwritten") else: print("ONNX simplify: check failed, kept unsimplified version") except Exception as e: print(f"ONNX simplify skipped: {e}, kept unsimplified version") # --------------------------------------------------------------------------- # Export entry-points # --------------------------------------------------------------------------- def export_image_encoder(model, output_path, num_classes=80, image_size=640): """Export Image Encoder + Neck + Head to ONNX.""" wrapper = ImageHeadWrapper(model) wrapper.eval() device = next(model.parameters()).device # Dummy inputs image = torch.randn(1, 3, image_size, image_size, device=device) text_feats = torch.randn(1, num_classes, 768, device=device) # Dry-run to verify shapes with torch.no_grad(): outputs = wrapper(image, text_feats) print("Image + Head output shapes:") for i, o in enumerate(outputs): print(f" output_{i}: {list(o.shape)}") output_names = [ "cls_scores_s8", "bbox_preds_s8", "cls_scores_s16", "bbox_preds_s16", "cls_scores_s32", "bbox_preds_s32", ] # Export to buffer → check → simplify → save with BytesIO() as f: torch.onnx.export( wrapper, (image, text_feats), f, input_names=["images", "text_features"], output_names=output_names, opset_version=14, do_constant_folding=True, ) f.seek(0) onnx_model = onnx.load(f) onnx.checker.check_model(onnx_model) _onnx_simplify(onnx_model, output_path) print(f"Exported: {output_path}") def export_text_encoder(model, output_path, num_texts=4, max_seq_len=32): """Export Text Encoder (XLM-RoBERTa + projection) to ONNX. The seq_len dimension is FIXED at export time. At inference, tokenized inputs must be padded to exactly ``max_seq_len`` tokens (the tokenizer does this automatically when ``padding='max_length'`` is set). Parameters ---------- max_seq_len : int Fixed token length. 32 is sufficient for typical Chinese class names. """ text_model = model.backbone.text_model wrapper = TextEncoderWrapper(text_model) wrapper.eval() device = next(text_model.parameters()).device # Dummy inputs — fixed seq_len, dynamic num_texts input_ids = torch.ones(num_texts, max_seq_len, dtype=torch.int64, device=device) attention_mask = torch.ones(num_texts, max_seq_len, dtype=torch.int64, device=device) with torch.no_grad(): out = wrapper(input_ids, attention_mask) print(f"Text Encoder output shape: {list(out.shape)}") # Export to buffer → check → simplify → save with BytesIO() as f: torch.onnx.export( wrapper, (input_ids, attention_mask), f, input_names=["input_ids", "attention_mask"], output_names=["text_features"], opset_version=17, do_constant_folding=True, ) f.seek(0) onnx_model = onnx.load(f) onnx.checker.check_model(onnx_model) _onnx_simplify(onnx_model, output_path) print(f"Exported: {output_path}") # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def parse_args(): parser = argparse.ArgumentParser(description="Export WeDetect to ONNX") parser.add_argument("--config", required=True, help="Config file path") parser.add_argument("--checkpoint", required=True, help="Checkpoint file path") parser.add_argument("--device", default="cuda:0") parser.add_argument("--num-classes", type=int, default=4, help="Number of classes for dummy text features (default 4).") parser.add_argument("--image-size", type=int, default=640) #此处演示导出检测4个类别的模型 parser.add_argument("--num-texts", type=int, default=4, help="Number of dummy texts for text-encoder trace.") #可以根据实际检测类别prompt计算选取最大seq_len parser.add_argument("--max-seq-len", type=int, default=32, help="Fixed token length for text-encoder ONNX. " "Inference inputs must be padded to this length.") parser.add_argument("--output-dir", default=".", help="Directory for the exported .onnx files.") return parser.parse_args() if __name__ == "__main__": args = parse_args() cfg = Config.fromfile(args.config) cfg.work_dir = osp.join("./work_dirs", osp.splitext(osp.basename(args.config))[0]) print(f"Loading model from {args.checkpoint} ...") model = init_detector(cfg, checkpoint=args.checkpoint, device=args.device, palette=['red']) model.eval() # ---- Image + Head ---- export_image_encoder( model, osp.join(args.output_dir, "wedetect_image_encoder.onnx"), num_classes=args.num_classes, image_size=args.image_size, ) # ---- Text Encoder ---- export_text_encoder( model, osp.join(args.output_dir, "wedetect_text_encoder.onnx"), num_texts=args.num_texts, max_seq_len=args.max_seq_len, ) print("Done.")