新增onnx导出脚本

export_onnx.py

@@ -0,0 +1,60 @@
+import torch
+import torchvision
+
+import onnx
+
+
+def get_args_parser(add_help=True):
+    import argparse
+
+    parser = argparse.ArgumentParser(description="model weight transport to onnx", add_help=add_help)
+
+    parser.add_argument("--model", default="resnet18", type=str, help="model name")
+    parser.add_argument(
+        "--num_classes", default=10, type=int, help="number of classes"
+    )
+    parser.add_argument(
+        "--input_size", default=224, type=int, help="input size"
+    )
+    parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load")
+    parser.add_argument("--save_path", default=None, type=str, help="onnx file save path")
+    return parser
+
+
+def export_onnx(args):
+    # 加载模型
+    model = torchvision.models.get_model(args.model, weights=None, num_classes=args.num_classes)
+    model.eval()  # 切换到评估模式
+
+    # 加载权重
+    checkpoint = torch.load(args.weights, map_location='cpu')
+    model.load_state_dict(checkpoint["model"])
+
+    # 定义一个随机输入张量，尺寸应符合模型输入要求
+    # 通常，ImageNet预训练模型的输入尺寸是 (batch_size, 3, 224, 224)
+    dummy_input = torch.randn(1, 3, args.input_size, args.input_size)
+
+    # 导出模型到 ONNX 格式
+    with torch.no_grad():
+        torch.onnx.export(
+            model,
+            dummy_input,
+            args.save_path,
+            export_params=True,
+            opset_version=11,
+            do_constant_folding=False,  # 确保这个参数一定是False，如果不为False，导出的onnx权重与原始权重数值不一致，导致白盒水印提取失败
+            input_names=["input"],
+            output_names=["output"],
+            dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
+        )
+
+    print("模型成功导出为 ONNX 格式！")
+
+    # Checks
+    model_onnx = onnx.load(args.save_path)  # load onnx model
+    onnx.checker.check_model(model_onnx)  # check onnx model
+
+
+if __name__ == '__main__':
+    args = get_args_parser().parse_args()
+    export_onnx(args)