oneflow-inc / oneflow_convert Goto Github PK

convert_to_onnx_and_check(t5_graph,
                  external_data=False, 
                  opset=None, 
                  flow_weight_dir=None, 
                  onnx_model_path="./", 
                  dynamic_batch_size=False)

convert_to_onnx_and_check(t5_graph,
                  external_data=False, 
                  opset=11, 
                  flow_weight_dir=None, 
                  onnx_model_path="./", 
                  dynamic_batch_size=False)

"""
Copyright 2020 The OneFlow Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import tempfile
import oneflow as flow
from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check

class Conv2d(flow.nn.Module):
    def __init__(self) -> None:
        super(Conv2d, self).__init__()
        self.conv = flow.nn.Conv2d(3, 16, 3)

    def forward(self, x: flow.Tensor) -> flow.Tensor:
        return self.conv(x)

conv_module = Conv2d()
class Conv2dOpGraph(flow.nn.Graph):
    def __init__(self):
        super().__init__()
        self.m = conv_module

    def build(self, x):
        out = self.m(x)
        return out


def test_conv2d():
    
    conv_graph = Conv2dOpGraph()
    conv_graph._compile(flow.randn(1, 3, 224, 224))

    with tempfile.TemporaryDirectory() as tmpdirname:
        flow.save(conv_module.state_dict(), tmpdirname)
        convert_to_onnx_and_check(conv_graph, flow_weight_dir=tmpdirname, onnx_model_path="/tmp", opset=14)

test_conv2d()

from oneflow_onnx.oneflow2onnx.util import export_onnx_model

export_onnx_model(graph,
                  external_data=False, 
                  opset=None, 
                  flow_weight_dir=None, 
                  onnx_model_path="/tmp", 
                  dynamic_batch_size=False)

import tempfile
import oneflow as flow
from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check

class logicalSliceAssign(flow.nn.Module):
    def __init__(self) -> None:
        super(logicalSliceAssign, self).__init__()
    
    def forward(self, x: flow.Tensor) -> flow.Tensor:
        x[:, 0 : 2] += x
        return x

logical_slice = logicalSliceAssign()
class logicalSliceOpGraph(flow.nn.Graph):
    def __init__(self):
        super().__init__()
        self.m = logical_slice

    def build(self, x):
        out = self.m(x)
        return out

def test_logical_slice():
    logical_slice_graph = logicalSliceOpGraph()
    logical_slice_graph._compile(flow.randn(1, 2, 1, 1))
    print(logical_slice_graph._ops_repr)

    with tempfile.TemporaryDirectory() as tmpdirname:
        flow.save(slice.state_dict(), tmpdirname)
        convert_to_onnx_and_check(logical_slice_graph, flow_weight_dir=tmpdirname, onnx_model_path="/tmp")

test_logical_slice()

import tempfile
import oneflow as flow
from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
from flowvision.models.face_recognition import iresnet50

model = iresnet50().to("cuda")

class ModelGraph(flow.nn.Graph):
    def __init__(self, model):
        super().__init__()
        self.config.allow_fuse_add_to_output(True)
        self.backbone = model

    def build(self, x):
        x = x.to("cuda")
        out = self.backbone(x)
        return out

model.eval()
model_graph = ModelGraph(model)
model_graph._compile(flow.randn(1, 3, 112, 112))


with tempfile.TemporaryDirectory() as tmpdirname:
    flow.save(model.state_dict(), tmpdirname)
    convert_to_onnx_and_check(
        model_graph, flow_weight_dir=tmpdirname, onnx_model_path="./", print_outlier=True)

# 假设class AlexNet存在
def alexnet(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> AlexNet:
    model = AlexNet(**kwargs)
    return model


def test_alexnet():
    alexnet_module = alexnet()
    alexnet_module.eval().to("cuda")
	
    # Graph添加到oneflow_onnx/oneflow2onnx/util.py文件中
    alexnet_graph = Graph(alexnet_module)
	
    # 这里把job换成了graph
    convert_to_onnx_and_check(
        alexnet_graph, 
        flow_weight_dir="./examples/oneflow2onnx/models/alexnet_oneflow_model", 
        onnx_model_path="/tmp"
    )

class Graph(flow.nn.Graph):
    def __init__(self, module):
        super().__init__()
        self.m = module

    def build(self, x):
        out = self.m(x)
        return out

shapes = {}
dtypes = {}
graph_str = repr(graph)
# print(graph_str)
size_where = 2
if "cuda" in graph_str:
    size_where = 3

p_size = re.compile(r"size=\(.*?\)", re.S)
p_type = re.compile(r"dtype=.*?,", re.S)
types = ["INPUT", "PARAMETER", "BUFFER", "OUTPUT"]

for t in types:
    data = re.finditer(t+":.*", graph_str)
    for i in data:
        attrs = i.group().split(":")
        size_str = re.findall(p_size, attrs[size_where])
        type_str = re.findall(p_type, attrs[size_where])
        assert size_str != [] or type_str != [], \
        "size should not be None, please check your inputs dtype"

        size_attr = size_str[0].replace("size=", "")
        type_attr = type_str[0].replace("dtype=", "")
        if size_attr[-2] == ",":
            size_attr = size_attr.replace(",", "")
        if type_attr[-1] == ",":
            type_str = type_attr.replace(",", "")

        data_size = tuple(map(int, size_attr[1:-1].split(", ")))
        node_name = attrs[1]
        shapes[node_name] = data_size
        dtypes[node_name] = STR_TO_FLOW[type_str]
            

def parse_attr(attr):
    # Parse node_attr
    attrs = {}
    for a in attr:
        attr_str = str(attr[a])

        if attr_str[0:7] == "at_list":
            attr_str_ = attr_str.split(" ")[0]

            if attr_str_ == "at_list_float":
                attrs[a] = tuple(attr[a].at_list_float.val)
            elif attr_str_ == "at_list_int32":
                attrs[a] = tuple(attr[a].at_list_int32.val)
            elif attr_str_ == "at_list_int64":
                attrs[a] = tuple(attr[a].at_list_int64.val)

        elif attr_str.split(":")[0] == "at_string":
            attrs[a] = attr[a].at_string

        elif attr_str.split(" ")[0] == "at_shape":
            attrs[a] = tuple(list(attr[a].at_shape.dim))

        else:
            attr_str_ = attr_str.split(":")[0]
            if attr_str_ == "at_bool":
                attrs[a] = attr[a].at_bool
            elif attr_str_ == "at_double":
                attrs[a] = attr[a].at_double
            elif attr_str_ == "at_float":
                attrs[a] = attr[a].at_float
            elif attr_str_ == "at_int32":
                attrs[a] = attr[a].at_int32
            elif attr_str_ == "at_int64":
                attrs[a] = attr[a].at_int64

    return attrs