diff --git a/examples/oneflow2onnx/models/CPU/test_free_eager_tensor.py b/examples/oneflow2onnx/models/CPU/test_free_eager_tensor.py new file mode 100644 index 0000000..81480d1 --- /dev/null +++ b/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) diff --git a/examples/oneflow2onnx/models/GPU/test_free_eager_tensor.py b/examples/oneflow2onnx/models/GPU/test_free_eager_tensor.py new file mode 100644 index 0000000..46710cc --- /dev/null +++ b/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") diff --git a/oneflow_onnx/onnx_wrapper.py b/oneflow_onnx/onnx_wrapper.py index aab1db8..6db8146 100644 --- a/oneflow_onnx/onnx_wrapper.py +++ b/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 = (