import argparse import os from dataclasses import dataclass from typing import Dict, Tuple import torch from groundingdino.models import build_model from groundingdino.util.slconfig import SLConfig from groundingdino.util.utils import clean_state_dict from groundingdino.util import get_tokenlizer from groundingdino.models.GroundingDINO.bertwarper import ( generate_masks_with_special_tokens_and_transfer_map, ) @dataclass class TextInputs: input_ids: torch.Tensor token_type_ids: torch.Tensor attention_mask: torch.Tensor position_ids: torch.Tensor text_self_attention_masks: torch.Tensor class GroundingDINOOnnxWrapper(torch.nn.Module): """ ONNX 导出用 wrapper： - 让 ONNX 输入完全是 Tensor（image + tokenized） - 复用 GroundingDINO 原始 backbone/transformer/head - 跳过 Python tokenizer/文本mask生成（这些在导出前就做好，作为输入喂进来） """ def __init__(self, model): super().__init__() self.model = model def forward( self, image: torch.Tensor, # [B,3,H,W] float32 input_ids: torch.Tensor, # [B,S] int64 token_type_ids: torch.Tensor, # [B,S] int64 attention_mask: torch.Tensor, # [B,S] int64 (用于 text_token_mask) position_ids: torch.Tensor, # [B,S] int64 text_self_attention_masks: torch.Tensor, # [B,S,S] bool/int64 (用于 sub_sentence_present) ) -> Tuple[torch.Tensor, torch.Tensor]: # ---- 文本编码（等价于 GroundingDINO.forward 内部实现）---- if self.model.sub_sentence_present: tokenized_for_encoder: Dict[str, torch.Tensor] = { "input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": text_self_attention_masks, "position_ids": position_ids, } else: tokenized_for_encoder = { "input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask, } bert_output = self.model.bert(**tokenized_for_encoder) encoded_text = self.model.feat_map(bert_output["last_hidden_state"]) text_token_mask = attention_mask.to(torch.bool) text_dict = { "encoded_text": encoded_text, "text_token_mask": text_token_mask, "position_ids": position_ids, "text_self_attention_masks": text_self_attention_masks.to(torch.bool), } # ---- 视觉编码 + transformer（基本照抄 GroundingDINO.forward）---- if isinstance(image, (list, torch.Tensor)): from groundingdino.util.misc import nested_tensor_from_tensor_list samples = nested_tensor_from_tensor_list(image) else: samples = image self.model.set_image_tensor(samples) import torch.nn.functional as F from groundingdino.util.misc import NestedTensor srcs = [] masks = [] for l, feat in enumerate(self.model.features): src, mask = feat.decompose() srcs.append(self.model.input_proj[l](src)) masks.append(mask) if self.model.num_feature_levels > len(srcs): _len_srcs = len(srcs) for l in range(_len_srcs, self.model.num_feature_levels): if l == _len_srcs: src = self.model.input_proj[l](self.model.features[-1].tensors) else: src = self.model.input_proj[l](srcs[-1]) m = samples.mask mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0] pos_l = self.model.backbone[1](NestedTensor(src, mask)).to(src.dtype) srcs.append(src) masks.append(mask) self.model.poss.append(pos_l) hs, reference, _, _, _ = self.model.transformer( srcs, masks, None, self.model.poss, None, None, text_dict ) from groundingdino.util.misc import inverse_sigmoid outputs_coord_list = [] for layer_ref_sig, layer_bbox_embed, layer_hs in zip( reference[:-1], self.model.bbox_embed, hs ): layer_delta_unsig = layer_bbox_embed(layer_hs) layer_outputs_unsig = layer_delta_unsig + inverse_sigmoid(layer_ref_sig) layer_outputs_unsig = layer_outputs_unsig.sigmoid() outputs_coord_list.append(layer_outputs_unsig) outputs_coord_list = torch.stack(outputs_coord_list) outputs_class = torch.stack( [ layer_cls_embed(layer_hs, text_dict) for layer_cls_embed, layer_hs in zip(self.model.class_embed, hs) ] ) pred_logits = outputs_class[-1] pred_boxes = outputs_coord_list[-1] self.model.unset_image_tensor() return pred_logits, pred_boxes def load_torch_model(model_config_path: str, model_checkpoint_path: str, device: str): args = SLConfig.fromfile(model_config_path) args.device = device model = build_model(args) checkpoint = torch.load(model_checkpoint_path, map_location="cpu") model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) model.eval() model.to(device) # ONNX trace 不兼容 gradient checkpoint，导出时统一关闭 for m in model.modules(): if hasattr(m, "use_checkpoint"): try: setattr(m, "use_checkpoint", False) except Exception: pass return model, args def build_text_inputs( tokenizer, caption: str, device: str, max_text_len: int, special_token_ids, ) -> TextInputs: caption = caption.lower().strip() if not caption.endswith("."): caption = caption + "." tokenized = tokenizer([caption], padding="longest", return_tensors="pt") tokenized = {k: v.to(device) for k, v in tokenized.items()} text_self_attention_masks, position_ids, _ = generate_masks_with_special_tokens_and_transfer_map( tokenized, special_token_ids, tokenizer ) # 按 max_text_len 裁剪（与模型 forward 行为一致） if text_self_attention_masks.shape[1] > max_text_len: s = max_text_len text_self_attention_masks = text_self_attention_masks[:, :s, :s] position_ids = position_ids[:, :s] tokenized["input_ids"] = tokenized["input_ids"][:, :s] tokenized["attention_mask"] = tokenized["attention_mask"][:, :s] tokenized["token_type_ids"] = tokenized["token_type_ids"][:, :s] return TextInputs( input_ids=tokenized["input_ids"].to(torch.int64), token_type_ids=tokenized["token_type_ids"].to(torch.int64), attention_mask=tokenized["attention_mask"].to(torch.int64), position_ids=position_ids.to(torch.int64), text_self_attention_masks=text_self_attention_masks, ) def main(): parser = argparse.ArgumentParser("Export GroundingDINO to ONNX", add_help=True) parser.add_argument("--config_file", "-c", type=str, required=True) parser.add_argument("--checkpoint_path", "-p", type=str, required=True) parser.add_argument("--output_onnx", "-o", type=str, required=True, help="输出 onnx 路径") parser.add_argument("--text_prompt", "-t", type=str, required=True, help="用于构建 dummy 文本输入（影响 seq_len）") parser.add_argument("--opset", type=int, default=17) parser.add_argument("--cpu-only", action="store_true") parser.add_argument("--dynamic", action="store_true", help="启用动态 H/W 与 seq_len 轴（更通用但可能更慢/更难优化）") parser.add_argument("--simplify", action="store_true", help="尝试用 onnxsim 简化（需要安装 onnxsim）") parser.add_argument("--image_hw", type=str, default="800,1333", help="dummy image H,W（默认与 transform 常见输出一致）") args = parser.parse_args() device = "cpu" if args.cpu_only else ("cuda" if torch.cuda.is_available() else "cpu") model, cfg = load_torch_model(args.config_file, args.checkpoint_path, device=device) tokenizer = get_tokenlizer.get_tokenlizer(cfg.text_encoder_type) special_token_ids = tokenizer.convert_tokens_to_ids(["[CLS]", "[SEP]", ".", "?"]) text_inputs = build_text_inputs( tokenizer=tokenizer, caption=args.text_prompt, device=device, max_text_len=getattr(cfg, "max_text_len", 256), special_token_ids=special_token_ids, ) h_str, w_str = args.image_hw.split(",") H, W = int(h_str), int(w_str) dummy_image = torch.randn(1, 3, H, W, device=device, dtype=torch.float32) wrapper = GroundingDINOOnnxWrapper(model) wrapper.eval() os.makedirs(os.path.dirname(os.path.abspath(args.output_onnx)) or ".", exist_ok=True) input_names = [ "image", "input_ids", "token_type_ids", "attention_mask", "position_ids", "text_self_attention_masks", ] output_names = ["pred_logits", "pred_boxes"] dynamic_axes = None if args.dynamic: dynamic_axes = { "image": {0: "batch", 2: "height", 3: "width"}, "input_ids": {0: "batch", 1: "seq"}, "token_type_ids": {0: "batch", 1: "seq"}, "attention_mask": {0: "batch", 1: "seq"}, "position_ids": {0: "batch", 1: "seq"}, "text_self_attention_masks": {0: "batch", 1: "seq", 2: "seq"}, "pred_logits": {0: "batch"}, "pred_boxes": {0: "batch"}, } with torch.no_grad(): torch.onnx.export( wrapper, ( dummy_image, text_inputs.input_ids, text_inputs.token_type_ids, text_inputs.attention_mask, text_inputs.position_ids, text_inputs.text_self_attention_masks, ), args.output_onnx, opset_version=args.opset, do_constant_folding=True, input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, ) if args.simplify: try: import onnx from onnxsim import simplify onnx_model = onnx.load(args.output_onnx) simplified_model, ok = simplify(onnx_model) if ok: onnx.save(simplified_model, args.output_onnx) print(f"✅ onnxsim 简化完成: {args.output_onnx}") else: print("⚠️ onnxsim 简化失败（模型未修改）") except Exception as e: print(f"⚠️ onnxsim 简化跳过: {e}") print(f"✅ 导出完成: {args.output_onnx}") if __name__ == "__main__": main()