support free eager tensor

examples/oneflow2onnx/models/CPU/test_free_eager_tensor.py

@@ -0,0 +1,31 @@
+import oneflow as flow
+import oneflow.nn as nn
+from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
+
+
+class ConvModule(nn.Module):
+    def __init__(self):
+        super(ConvModule, self).__init__()
+        self.zero = flow.zeros(1, 2, 3, 3)
+        self.conv = nn.Conv2d(1, 2, 1, 1)
+
+    def forward(self, x):
+        return self.conv(x) + self.zero
+
+
+m = ConvModule()
+
+
+class YOLOGraph(flow.nn.Graph):
+    def __init__(self, m):
+        super().__init__()
+        self.model = m
+
+    def build(self, x):
+        return self.model(x)
+
+
+yolo_graph = YOLOGraph(m)
+yolo_graph._compile(flow.randn(1, 1, 3, 3))
+
+convert_to_onnx_and_check(yolo_graph, onnx_model_path="/tmp", print_outlier=True)

examples/oneflow2onnx/models/GPU/test_free_eager_tensor.py

@@ -0,0 +1,32 @@
+import oneflow as flow
+import oneflow.nn as nn
+from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
+
+
+class ConvModule(nn.Module):
+    def __init__(self):
+        super(ConvModule, self).__init__()
+        self.zero = flow.zeros(1, 2, 3, 3).to("cuda")
+        self.conv = nn.Conv2d(1, 2, 1, 1)
+
+    def forward(self, x):
+        return self.conv(x) + self.zero
+
+
+m = ConvModule()
+m.to("cuda")
+
+
+class YOLOGraph(flow.nn.Graph):
+    def __init__(self, m):
+        super().__init__()
+        self.model = m
+
+    def build(self, x):
+        return self.model(x)
+
+
+yolo_graph = YOLOGraph(m)
+yolo_graph._compile(flow.randn(1, 1, 3, 3).to("cuda"))
+
+convert_to_onnx_and_check(yolo_graph, onnx_model_path="/tmp", device="gpu")

oneflow_onnx/onnx_wrapper.py

@@ -636,7 +636,7 @@ def UpdateNodeShapeDtype(self, node, override=False):
                 if tensor_value.size * tensor_value.itemsize > 128 * 1024 * 1024:  # 128 MB
                     tmp_dir = tempfile.TemporaryDirectory()
                     tmp_dirs.append(tmp_dir)
-                    tensor = util.TensorProtoFromNumpy(tensor_value, name=inp.output_tensor_names[0], external_data=True, export_path=tmp_dir.name)
+                    tensor = util.TensorProtoFromNumpy(tensor_value, name=inp.output_tensor_names[0], external_data=True, export_path=tmp_dir.name,)
                 else:
                     tensor = util.TensorProtoFromNumpy(tensor_value, name=inp.output_tensor_names[0])
                 initializers.append(tensor)
@@ -783,8 +783,20 @@ def get_saved_tensor(self, node):
         # TODO(daquexian): node.output_tensor_names[0] is "node_name/output_name", so this pathjoin doesn't work
         # on windows (where path separator is "\")
         key = ".".join(node.output_tensor_names[0].split(".")[1:])
+        # For Free Eager Tensor
+        if key == "":
+            key = node.output_tensor_names[0]
         if len(list(self._param_dict.keys())) > 1:
-            tensor_value = self._param_dict[key[: key.rfind("out") - 1]].numpy().reshape(self.get_shape(tensor_name)).astype(dtype=util.Onnx2NumpyDtype(self.get_dtype(tensor_name)))
+            try:
+                tensor_value = self._param_dict[key[: key.rfind("out") - 1]].numpy().reshape(self.get_shape(tensor_name)).astype(dtype=util.Onnx2NumpyDtype(self.get_dtype(tensor_name)))
+            except:
+                # For Eager Free Tensor
+                tensor_value = (
+                    self._param_dict[list(self._param_dict.keys())[0]][key[: key.rfind("out") - 1]]
+                    .numpy()
+                    .reshape(self.get_shape(tensor_name))
+                    .astype(dtype=util.Onnx2NumpyDtype(self.get_dtype(tensor_name)))
+                )
         else:
             try:
                 tensor_value = (