export to onnx

export_to_torchscript.py

@@ -84,10 +84,20 @@ def ctranspath():
 
 model = ctranspath()
 model.head = torch.nn.Identity()
-td = torch.load(r"./ctranspath.pth")
+td = torch.load("ctranspath.pth")
 model.load_state_dict(td["model"], strict=True)
 
-
 jitted = torch.jit.script(model)
-
 torch.jit.save(jitted, "torchscript_model.pt")
+
+torch.onnx.export(
+    model,
+    args=torch.ones(1, 3, 224, 224),
+    f="model.onnx",
+    input_names=["image"],
+    output_names=["embedding"],
+    dynamic_axes={
+        "input": {0: "batch_size"},
+        "output": {0: "batch_size"},
+    },
+)

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eec3f58a2b57142c48411467858847f75b2726f7b22db877428fcf06aaac2958
+size 112312109