OneFlow是如何和ONNX交互的

本篇內容主要講解“OneFlow是如何和ONNX交互的”，感興趣的朋友不妨來看看。本文介紹的方法操作簡單快捷，實用性強。下面就讓小編來帶大家學習“OneFlow是如何和ONNX交互的”吧!

0x0. 介紹

在開始閱讀本篇文章之前，如果你對ONNX不是很了解介意先閱讀我之前寫的這幾篇介紹ONNX文章：

• ONNX初探

• ONNX 再探

• onnx2pytorch和onnx-simplifier新版介紹

以及大老師的：

onnx simplifier 和 optimizer

然后，這篇文章不會繼續探索ONNX本身的東西，而是聊聊另外一個有趣的話題，即深度學習框架是如何和ONNX進行交互的？我最近配合大老師基于OneFlow深度學習框架做了一些和ONNX有關的工作，感覺自己對OneFlow和ONNX的交互過程也算熟悉一些了。因此，在這篇文章我將分享OneFlow和ONNX交互的具體實現思路以及介紹oneflow-onnx這個開源工具的一些特性。讓讀者了解OneFlow的模型是如何轉換為ONNX模型，以及ONNX模型是如何轉回OneFlow的模型（X2OneFlow）的。個人認為OneFlow目前和ONNX交互的做法是比較優雅并且具有較好擴展性的，因此我們將這項工作轉換成了開源成果并分享實現思路，github地址為：https://github.com/Oneflow-Inc/oneflow_convert_tools。這個工具作為OneFlow 生態系統的一部分會被我們持續維護，同時這個工具也被我們制作成了一個wheel包，感興趣的用戶只需要pip安裝oneflow-onnx即可快速體驗。在下面的第二節以及工程的README也有詳細的安裝步驟。

oneflow-onnx工具包含兩個功能，一個是將OneFlow導出ONNX，另外一個是將各個訓練框架導出的ONNX模型轉換為OneFlow的模型。本工程已經適配了TensorFlow/Pytorch/PaddlePaddle框架的預訓練模型通過導出ONNX轉換為OneFlow（我們將這一功能叫作X2OneFlow）。更多使用示例以及相關文檔和源碼均可以在開源https://github.com/Oneflow-Inc/oneflow_convert_tools工程中獲得。

0x1. 算子支持和模型支持

OneFlow2ONNX

OneFlow2ONNX 支持的OP列表

目前OneFlow2ONNX 支持60+的ONNX OP，我們在下面的列表中列出了目前OneFlow支持導出的全部ONNX OP

OneFlow2ONNX 模型測試庫

目前OneFlow2ONNX 支持60+的ONNX OP，我們在下面的模型列表中測試了OneFlow2ONNX的轉換。

X2OneFlow

X2OneFlow 支持的OP列表

目前X2OneFlow 支持40+的ONNX OP，30+的Tensorflow/Pytorch/PaddlePaddle OP，覆蓋了大部分CV分類模型常用的操作。OP的單元測試代碼會逐漸移步到工程的examples目錄下，并支持更多的OP。

ONNX

TensorFlow

• 分組卷積存在問題，已給TensorFlow2ONNX團隊PR。

Pytorch

PaddlePaddle

相關issue：

• https://github.com/PaddlePaddle/Paddle2ONNX/issues/221

• https://github.com/PaddlePaddle/Paddle2ONNX/issues/220

X2OneFlow模型測試庫

目前X2OneFlow 支持40+的ONNX OP，30+的Tensorflow/Pytorch/PaddlePaddle OP，覆蓋了大部分CV分類模型常用的操作。我們在如下模型列表中測試了X2OneFlow的轉換。

Pytorch

TensorFlow

PaddlePaddle

• 模型的測試代碼均可以在工程的examples中找到。

0x2. 快速體驗

用戶環境配置

python>=3.5
onnx>=1.8.0
onnx-simplifier>=0.3.3
onnxoptimizer>=0.2.5
onnxruntime>=1.6.0
oneflow>=0.3.4

如果你想使用X2OneFlow（X代表TensorFlow/Pytorch/PaddlePaddle）需要安裝對應的深度學習框架，需要安裝對應的深度學習框架，依賴如下：

pytorch>=1.7.0
paddlepaddle>=2.0.0
tensorflow>=2.0.0

安裝

安裝方式1

pip install oneflow_onnx

安裝方式2

git clone https://github.com/Oneflow-Inc/oneflow_convert_toolscd oneflow_onnx
python3 setup.py install

使用方法見工程的samples下的示例。

0x3. OneFlow-ONNX思路分享

我們將在這一節分享一下OneFlow的模型是如何被轉換為ONNX的，這里我們以將OneFlow定義的AlexNet導出ONNX模型為例來分析源碼。首先我們https://github.com/Oneflow-Inc/oneflow_convert_tools/blob/main/examples/oneflow2onnx/test_alexnet.py#L133進到這里，可以看到下面調用代碼：

def test_alexnet():    @flow.global_function()    def alexnet_eval_job(x: tp.Numpy.Placeholder((1, 227, 227, 3))):        return alexnet(x, None, False)

    convert_to_onnx_and_check(alexnet_eval_job, flow_weight_dir=None, onnx_model_path="/tmp")

這里通過flow.global_function()定義了一個預測用于eval的AlexNet job，網絡的完整定義可以通過上面的鏈接訪問，可以看到這里通過convert_to_onnx_and_check函數將OneFlow定義的AlexNet轉換為了ONNX模型，我們跟進這個函數，就來到了這里：https://github.com/Oneflow-Inc/oneflow_convert_tools/blob/main/oneflow_onnx/oneflow2onnx/util.py#L65-L73，代碼為：

 while not os.path.exists(os.path.join(flow_weight_dir, "snapshot_done")):
     pass onnx_model_dir = onnx_model_path
 onnx_model_path = os.path.join(onnx_model_dir, "model.onnx")
 flow.onnx.export(
     job_func,
     flow_weight_dir,
     onnx_model_path,
     opset=opset,
     external_data=external_data,
 )

可以看到完成ONNX模型轉換的核心函數就是這個flow.onnx.export函數，我們繼續跳轉到這個函數https://github.com/Oneflow-Inc/oneflow_convert_tools/blob/main/oneflow_onnx/oneflow2onnx/flow2onnx.py#L229-L281，代碼如下：

def Export(
    job_func: Callable,
    model_save_dir: Text,
    onnx_filename: Text,
    continue_on_error: bool = False,
    opset: Optional[int] = None,
    extra_opset: Optional[int] = None,
    shape_override: Optional[Dict[Text, List[int]]] = None,
    external_data: bool = False,
):    r"""Export a oneflow model into ONNX format.
    Args:
        job_func: OneFlow的作業函數
        model_save_dir: 包含OneFlow定義的模型權重的文件夾. 這個模型權重是用oneflow的check_point.save接口保存的。
        onnx_filename: 輸出ONNX模型文件名，字符串類型
        continue_on_error: 如果某個OP無法處理（即沒有映射），是否繼續
        opset: ONNX Opset版本號，默認為10
        extra_opset: 額外Opset的列表，例如自定義操作使用的Opset
        shape_override: 帶有輸入信息的字典，覆蓋OneFlow給定的輸入形狀
        external_data: 將權重另存為ONNX外部數據，通常是為了繞過protobuf的2GB文件大小限制。
    """    assert os.getenv("ENABLE_USER_OP") != "False"    # 確定模型的路徑是存在的    assert os.path.isdir(model_save_dir)
    # 通過c_api_util.GetJobSet()方法獲取當前的所有job    job_set = c_api_util.GetJobSet()
    # 我們要轉的模型被定義在job_func中，所以我們先記錄下它的名字    job_name = job_func.__name__
    # 編譯job_set，找到定義模型的job    for job in job_set.job:
        # TODO(OYY) Modify the interface before modifying it        if job.job_conf.job_name == job_name:
            # job找到了，可以開始進行下面的步驟，我們在外面詳細解釋            onnx_graph = ProcessFlowGraph(
                job,
                model_save_dir,
                continue_on_error=continue_on_error,
                opset=opset,
                extra_opset=extra_opset,
                shape_override=shape_override,
            )
            onnx_graph = optimizer.OptimizeGraph(onnx_graph)
            model_proto = onnx_graph.MakeModel(
                job_name, onnx_filename, external_data=external_data
            )
            with open(onnx_filename, "wb") as f:
                try:
                    f.write(model_proto.SerializeToString())
                except ValueError as e:
                    raise ValueError(
                        "Error occured when running model_proto.SerializeToString(). If the model is larger than 2GB, please specify external_data=True when calling flow.onnx.export. Original error message:\n{}".format(
                            e
                        )
                    )
            return    raise ValueError('Cannot find job "{}" in jobset'.format(job_name))

可以看到這個函數首先編譯了OneFlow中的job_set，然后找到了我們最開始定義AlexNet模型的那個job，然后就進入了ProcessFlowGraph函數，這個函數主要做了三件事情并最終獲得了初版的合法ONNX模型（初版的意思是還沒有經過優化以及填ONNX節點的權重），我們跟進這個函數，代碼如下。

def ProcessFlowGraph(
    flow_graph,
    model_save_dir,
    continue_on_error=False,
    opset=None,
    extra_opset=None,
    shape_override=None,
): # 這個函數用來獲取導出的ONNX的Opset Version，OneFlow里面最高為10    opset = util.FindOpset(opset)
    logger.info("Using opset <onnx, %s>", opset)
    # 判斷當前的ONNX版本是否支持上面的Opset Version    if opset > schemas.get_max_supported_opset_version():
        logger.warning(
            "Currently installed onnx package %s is too low to support opset %s, "            "please upgrade onnx package to avoid potential conversion issue.",
            util.get_onnx_version(),
            opset,
        )

    if shape_override is None:
        shape_override = {}
 # 用于將oneflow 的各個 node 轉換為 onnx node 的格式，保持 op 類型、輸入輸出和屬性值不變，這一步產生的還不是合法的 onnx 模型    (onnx_nodes, op_cnt, attr_cnt, dtypes, output_shapes,) = FlowToOnnxNaive(
        flow_graph, shape_override
    )
 # 構造一個 Graph 類，用于后續方便的修改 onnx 網絡結構    g = Graph(onnx_nodes, model_save_dir, output_shapes, dtypes, opset, extra_opset,)

    # create ops mapping for the desired opsets    ops_mapping = handler.flow_op.CreateMapping(g.opset, g.extra_opset)

    # some nodes may already copied into inner Graph, so remove them from main Graph.    TopologicalSort(g, continue_on_error)
 # FlowOnnxMapping 函數調用各個轉換函數（通過 @flow_op 注冊）逐個轉換 op，轉換后產生的是合法的 onnx 模型    mapped_op, unmapped_op, exceptions = FlowOnnxMapping(g, ops_mapping)
    if unmapped_op:
        logger.error("Unsupported ops: %s", unmapped_op)
    if exceptions and not continue_on_error:
        raise exceptions[0]

    # onnx requires topological sorting    TopologicalSort(g, continue_on_error)

    g.UpdateProto()

    logger.debug(
        "Summay Stats:\n"        "\toneflow ops: {}\n"        "\toneflow attr: {}\n"        "\tonnx mapped: {}\n"        "\tonnx unmapped: {}".format(op_cnt, attr_cnt, mapped_op, unmapped_op)
    )

    return g

FlowToOnnxNaive這個函數用于將oneflow 的各個 node 轉換為 onnx node 的格式，保持 op 類型、輸入輸出和屬性值不變，最后將轉換后的ONNX節點（這個地方這些ONNX節點還不是真正的合法ONNX節點，要后面執行一對一轉換之后才是合法的ONNX節點）全部返回。接下來利用這些ONNX節點來構造Graph類，方便后續對ONNX模型進行修改。Graph類的實現在https://github.com/Oneflow-Inc/oneflow_convert_tools/blob/18e041d92654cfc8b03e16c906c451a405c99fd2/oneflow_onnx/onnx_wrapper.py，這個文件主要是定義了onnx graph和node的wrapper，包含各種修改 onnx 圖結構的 api，這里復用了tensorflow-onnx項目的相關代碼。注意構造Graph類之后還并沒有構造ONNX模型，因為OneFlow的OP還沒有一對一的轉換為ONNX的OP。

接下來，我們調用handler.flow_op.CreateMapping(g.opset, g.extra_opset)這個函數，代碼實現如下：

def CreateMapping(max_onnx_opset_version, extra_opsets):        """Create the final mapping dictionary by stacking domains and opset versions.

        :param max_onnx_opset_version: The highest onnx opset the resulting graph may use.
        :param extra_opsets: Extra opsets the resulting graph may use.
        """        mapping = {constants.ONNX_DOMAIN: max_onnx_opset_version}
        if extra_opsets:
            for extra_opset in extra_opsets:
                mapping[extra_opset.domain] = extra_opset.version
        ops_mapping = {}
        for domain, opsets in flow_op.get_opsets().items():
            for target_opset, op_map in enumerate(opsets):
                print('='*100)
                print(target_opset)
                print(op_map)
                m = mapping.get(domain)
                if m:
                    if target_opset <= m and op_map:
                        ops_mapping.update(op_map)

        flow_op._MAPPING = ops_mapping
        return ops_mapping

這個函數做的事情就是將每個ONNX Opset版本號（也就是for循環中的domain）和(OneFlow OP和ONNX OP的mapper，這個mapper是如何獲得的請看后文)關聯起來并返回，我們打印一下target_opset和op_map就可以理解了。以AlexNet為例打印如下：

====================================================================================================0{}
====================================================================================================1{'add_n': (<bound method AddN.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.AddN'>>, 'Sum', {}), 'leaky_relu': (<bound method DirectOpSinceOpset1.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOpSinceOpset1'>>, 'LeakyRelu', {}), 'softplus': (<bound method DirectOpSinceOpset1.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOpSinceOpset1'>>, 'Softplus', {}), 'abs': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Abs', {}), 'ceil': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Ceil', {}), 'elu': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Elu', {}), 'exp': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Exp', {}), 'floor': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Floor', {}), 'log': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Log', {}), 'neg': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Neg', {}), 'sigmoid': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Sigmoid', {}), 'sigmoid_v2': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Sigmoid', {}), 'sqrt': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Sqrt', {}), 'tanh': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Tanh', {}), 'reciprocal': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Reciprocal', {}), 'relu': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Relu', {}), 'broadcast_maximum': (<bound method MinMaxOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.MinMaxOp'>>, 'Max', {}), 'broadcast_minimum': (<bound method MinMaxOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.MinMaxOp'>>, 'Min', {}), 'clip_by_scalar': (<bound method ClipByValueOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.ClipByValueOp'>>, 'Clip', {}), 'clip_by_scalar_min': (<bound method ClipByValueOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.ClipByValueOp'>>, 'Clip', {}), 'clip_by_scalar_max': (<bound method ClipByValueOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.ClipByValueOp'>>, 'Clip', {}), 'softmax': (<bound method Softmax.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Softmax'>>, 'Softmax', {}), 'square': (<bound method Square.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Square'>>, None, {}), 'rsqrt': (<bound method Rsqrt.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Rsqrt'>>, None, {}), 'squared_difference': (<bound method SquaredDifference.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.SquaredDifference'>>, None, {}), 'sign': (<bound method Sign.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Sign'>>, 'Sign', {}), 'matmul': (<bound method MatMul.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.MatMul'>>, 'MatMul', {}), 'batch_matmul': (<bound method MatMul.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.MatMul'>>, 'MatMul', {}), 'erf': (<bound method Erf.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Erf'>>, 'Erf', {}), 'logical_not': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Not', {}), 'broadcast_logical_or': (<bound method BroadcastOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Or', {}), 'broadcast_logical_and': (<bound method BroadcastOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'And', {}), 'input': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.misc.DirectOp'>>, None, {}), 'return': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.misc.DirectOp'>>, None, {}), 'variable': (<bound method DirectOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.misc.DirectOp'>>, None, {}), 'distribute_split': (<bound method BoxingOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.misc.BoxingOp'>>, 'Identity', {}), 'distribute_concat': (<bound method BoxingOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.misc.BoxingOp'>>, 'Identity', {}), 'distribute_clone': (<bound method BoxingOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.misc.BoxingOp'>>, 'Identity', {}), 'distribute_add': (<bound method BoxingOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.misc.BoxingOp'>>, 'Identity', {}), 'conv2d': (<bound method ConvOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.ConvOp'>>, None, {}), 'max_pool_2d': (<bound method PoolOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.PoolOp'>>, 'MaxPool', {}), 'avg_pool_2d': (<bound method PoolOp.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.PoolOp'>>, 'AveragePool', {}), 'reduce_prod': (<bound method ReduceOpBase.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ReduceOpBase'>>, 'ReduceProd', {}), 'reduce_sum': (<bound method ReduceOpBase.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ReduceOpBase'>>, 'ReduceSum', {}), 'reduce_min': (<bound method ReduceOpBase.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ReduceOpBase'>>, 'ReduceMin', {}), 'argmax': (<bound method ArgMax.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ArgMax'>>, 'ArgMax', {}), 'argmin': (<bound method ArgMax.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ArgMax'>>, 'ArgMin', {}), 'squeeze': (<bound method Squeeze.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Squeeze'>>, 'Squeeze', {}), 'transpose': (<bound method Transpose.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Transpose'>>, 'Transpose', {}), 'concat': (<bound method Concat.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Concat'>>, 'Concat', {}), 'identity': (<bound method Identity.Version_1 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Identity'>>, 'Identity', {})}
====================================================================================================
2
{'pad': (<bound method Pad.Version_2 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.Pad'>>, 'Pad', {})}
====================================================================================================
3
{}
====================================================================================================
4
{}
====================================================================================================
5
{'reshape': (<bound method Reshape.Version_5 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Reshape'>>, 'Reshape', {})}
====================================================================================================
6
{'broadcast_div': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Div', {}), 'scalar_div_by_tensor': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Div', {}), 'multiply': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Mul', {}), 'broadcast_mul': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Mul', {}), 'scalar_mul_by_tensor': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Mul', {}), 'broadcast_sub': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Sub', {}), 'scalar_sub_by_tensor': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Sub', {}), 'broadcast_add': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Add', {}), 'scalar_add_by_tensor': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Add', {}), 'scalar_add': (<bound method ScalarBinaryOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.ScalarBinaryOp'>>, 'Add', {}), 'scalar_mul': (<bound method ScalarBinaryOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.ScalarBinaryOp'>>, 'Mul', {}), 'bias_add': (<bound method BiasAdd.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BiasAdd'>>, 'Add', {}), 'abs': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Abs', {}), 'ceil': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Ceil', {}), 'elu': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Elu', {}), 'exp': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Exp', {}), 'floor': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Floor', {}), 'log': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Log', {}), 'neg': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Neg', {}), 'sigmoid': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Sigmoid', {}), 'sigmoid_v2': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Sigmoid', {}), 'sqrt': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Sqrt', {}), 'tanh': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Tanh', {}), 'reciprocal': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Reciprocal', {}), 'relu': (<bound method DirectOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.DirectOp'>>, 'Relu', {}), 'broadcast_logical_or': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'Or', {}), 'broadcast_logical_and': (<bound method BroadcastOp.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.BroadcastOp'>>, 'And', {}), 'normalization': (<bound method BatchNorm.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.BatchNorm'>>, None, {}), 'cast': (<bound method Cast.Version_6 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Cast'>>, 'Cast', {})}
====================================================================================================
7
{'acos': (<bound method TrigOpSinceOpset7.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset7'>>, 'Acos', {}), 'asin': (<bound method TrigOpSinceOpset7.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset7'>>, 'Asin', {}), 'atan': (<bound method TrigOpSinceOpset7.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset7'>>, 'Atan', {}), 'cos': (<bound method TrigOpSinceOpset7.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset7'>>, 'Cos', {}), 'sin': (<bound method TrigOpSinceOpset7.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset7'>>, 'Sin', {}), 'tan': (<bound method TrigOpSinceOpset7.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset7'>>, 'Tan', {}), 'broadcast_floor_mod': (<bound method FloorMod.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.FloorMod'>>, 'FloorMod', {}), 'broadcast_equal': (<bound method Equal.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Equal'>>, 'Equal', {}), 'broadcast_not_equal': (<bound method Equal.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Equal'>>, 'NotEqual', {}), 'broadcast_greater': (<bound method GreaterLess.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.GreaterLess'>>, 'Greater', {}), 'broadcast_less': (<bound method GreaterLess.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.GreaterLess'>>, 'Less', {}), 'broadcast_less_equal': (<bound method GreaterLessEqual.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.GreaterLessEqual'>>, 'Greater', {}), 'broadcast_greater_equal': (<bound method GreaterLessEqual.Version_7 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.GreaterLessEqual'>>, 'Less', {})}
====================================================================================================
8
{}
====================================================================================================
9
{'acosh': (<bound method TrigOpSinceOpset9.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset9'>>, 'Acosh', {}), 'asinh': (<bound method TrigOpSinceOpset9.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset9'>>, 'Asinh', {}), 'atanh': (<bound method TrigOpSinceOpset9.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset9'>>, 'Atanh', {}), 'cosh': (<bound method TrigOpSinceOpset9.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset9'>>, 'Cosh', {}), 'sinh': (<bound method TrigOpSinceOpset9.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.TrigOpSinceOpset9'>>, 'Sinh', {}), 'sign': (<bound method Sign.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Sign'>>, 'Sign', {}), 'erf': (<bound method Erf.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Erf'>>, 'Erf', {}), 'normalization': (<bound method BatchNorm.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.BatchNorm'>>, None, {}), 'cast': (<bound method Cast.Version_9 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Cast'>>, 'Cast', {})}
====================================================================================================
10
{'max_pool_2d': (<bound method PoolOp.Version_10 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.PoolOp'>>, 'MaxPool', {}), 'avg_pool_2d': (<bound method PoolOp.Version_10 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.PoolOp'>>, 'AveragePool', {}), 'min_max_observer': (<bound method MinMaxObserver.Version_10 of <class 'oneflow_onnx.oneflow2onnx.handlers.quantize.MinMaxObserver'>>, None, {}), 'moving_average_min_max_observer': (<bound method MovingAverageMinMaxObserver.Version_10 of <class 'oneflow_onnx.oneflow2onnx.handlers.quantize.MovingAverageMinMaxObserver'>>, None, {}), 'fake_quantization': (<bound method FakeQuantization.Version_10 of <class 'oneflow_onnx.oneflow2onnx.handlers.quantize.FakeQuantization'>>, 'QuantizeLinear', {})}
====================================================================================================
11
{'clip_by_scalar': (<bound method ClipByValueOp.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.ClipByValueOp'>>, 'Clip', {}), 'clip_by_scalar_min': (<bound method ClipByValueOp.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.ClipByValueOp'>>, 'Clip', {}), 'clip_by_scalar_max': (<bound method ClipByValueOp.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.ClipByValueOp'>>, 'Clip', {}), 'softmax': (<bound method Softmax.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Softmax'>>, 'Softmax', {}), 'round': (<bound method Round.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Round'>>, 'Round', {}), 'broadcast_equal': (<bound method Equal.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Equal'>>, 'Equal', {}), 'broadcast_not_equal': (<bound method Equal.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.math.Equal'>>, 'NotEqual', {}), 'conv2d': (<bound method ConvOp.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.ConvOp'>>, None, {}), 'max_pool_2d': (<bound method PoolOp.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.PoolOp'>>, 'MaxPool', {}), 'avg_pool_2d': (<bound method PoolOp.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.PoolOp'>>, 'AveragePool', {}), 'pad': (<bound method Pad.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.nn.Pad'>>, 'Pad', {}), 'reduce_prod': (<bound method ReduceOpBase.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ReduceOpBase'>>, 'ReduceProd', {}), 'reduce_sum': (<bound method ReduceOpBase.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ReduceOpBase'>>, 'ReduceSum', {}), 'reduce_min': (<bound method ReduceOpBase.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ReduceOpBase'>>, 'ReduceMin', {}), 'argmax': (<bound method ArgMax.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ArgMax'>>, 'ArgMax', {}), 'argmin': (<bound method ArgMax.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.reduce.ArgMax'>>, 'ArgMin', {}), 'squeeze': (<bound method Squeeze.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Squeeze'>>, 'Squeeze', {}), 'concat': (<bound method Concat.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.Concat'>>, 'Concat', {}), 'gather_nd': (<bound method GatherND.Version_11 of <class 'oneflow_onnx.oneflow2onnx.handlers.array.GatherND'>>, 'GatherND', {})}
====================================================================================================
12
{}
====================================================================================================
13
{'min_max_observer': (<bound method MinMaxObserver.Version_13 of <class 'oneflow_onnx.oneflow2onnx.handlers.quantize.MinMaxObserver'>>, None, {}), 'fake_quantization': (<bound method FakeQuantization.Version_13 of <class 'oneflow_onnx.oneflow2onnx.handlers.quantize.FakeQuantization'>>, 'QuantizeLinear', {})}

可以看到對于ONNX的每一個Opset Version的OP都對應了OneFlow實現的OP，需要特別注意的是這個OP Mapper過程在是在https://github.com/Oneflow-Inc/oneflow_convert_tools/tree/main/oneflow_onnx/oneflow2onnx/handlers這里完成的，只要安裝了oneflow-onnx這個包或者編譯了oneflow-onnx工程源碼，Python就會自動將OneFlow的OP和ONNX的OP進行映射，這是通過@flow_op(["avg_pool_2d"], onnx_op="AveragePool")裝飾器來實現的，flow_op裝飾器的具體實現在https://github.com/Oneflow-Inc/oneflow_convert_tools/blob/main/oneflow_onnx/oneflow2onnx/handler.py#L34這里。

完成了ONNX每個Opset版本的OP和OneFlow OP的mapper之后，我們需要對Graph里面的ONNX節點（注意現在的ONNX節點并不是合法的ONNX節點，因為還沒有執行一對一的轉換，只是復制了OneFlow OP的類型、輸入輸出和屬性值）先執行拓撲排序，然后再一對一的轉換。這個地方很有意思，為什么要進行拓撲排序呢？

我們首先需要了解一下拓撲序算法，拓撲排序要解決的問題是給一個圖的所有節點排序。

以下對拓撲排序的解釋引自oi.wiki。

我們可以拿大學選課的例子來描述這個過程，比如學習大學課程中有：單變量微積分，線性代數，離散數學概述，概率論與統計學概述，語言基礎，算法導論，機器學習。當我們想要學習 算法導論 的時候，就必須先學會 離散數學概述 和 概率論與統計學概述，不然在課堂就會聽的一臉懵逼。當然還有一個更加前的課程 單變量微積分。這些課程就相當于幾個頂點 , 頂點之間的有向邊  就相當于學習課程的順序。顯然拓撲排序不是那么的麻煩，不然你是如何選出合適的學習順序。下面將介紹如何將這個過程抽象出來，用算法來實現。

但是如果某一天排課的老師打瞌睡了，說想要學習 算法導論，還得先學 機器學習，而 機器學習 的前置課程又是 算法導論，然后你就一萬臉懵逼了，我到底應該先學哪一個？當然我們在這里不考慮什么同時學幾個課程的情況。在這里，算法導論 和 機器學習 間就出現了一個環，顯然你現在沒辦法弄清楚你需要學什么了，于是你也沒辦法進行拓撲排序了。因而如果有向圖中存在環路，那么我們就沒辦法進行 拓撲排序 了。

因此我們可以說 在一個 DAG（有向無環圖），我們將圖中的頂點以線性方式進行排序，使得對于任何的頂點  到  的有向邊 , 都可以有  在  的前面。

還有給定一個 DAG，如果從  到  有邊，則認為  依賴于 。如果  到  有路徑（ 可達 ），則稱  間接依賴于 。

拓撲排序的目標是將所有節點排序，使得排在前面的節點不能依賴于排在后面的節點。 偽代碼實現如下：

void TopologicalSort(Graph G){
 InitStack(S);
 for(i = 0;i < G.vexnum; i++){
  if(indegrdd[i]==0)
   Push(S, i);
 }
 
 int count =0;
 while(!Empty(S)){
  Pop(S,i);
  print[count++] = i;
  for(p = G.vertices[i].firstarc; p; p = p->nextarc){
   v = p->adjvex;
   if(!(--indegree[v]))
    Push(S, v);
  }
 }
 
 if(count < G.vexnum)
  return false;
 else  return true;
}

上面加粗的這句話即是拓撲排序的核心。一般深度學習模型也是一個DAG（有向無環圖），我們這里同樣使用了拓撲排序算法使得我們在一對一轉換OP時和真實的網絡結構是完全一致的。另外考慮到這里可能插入了一些新的節點如Identity可能會破壞原Graph的拓撲序，以及時刻需要判斷計算圖是否是一個完整合法的DAG，使用拓撲排序都是沒有壞處的。

完成拓撲排序之后我們就可以執行FlowOnnxMapping完成OneFlow OP和ONNX OP的一對一轉換了，代碼如下：

def FlowOnnxMapping(g, ops_mapping):    logger.debug("Mapping Oneflow node to ONNX node(s)")
    mapped_op = collections.Counter()
    unmapped_op = collections.Counter()
    exceptions = []

    ops = list(g.get_nodes())
    for node in ops:
        logger.debug("Process node: %s\n%s", node.name, node.summary)

        if node.skip_conversion:
            logger.debug("explicitly skip node " + node.name)
            continue        op = node.op_type
        map_info = ops_mapping.get(op)
        if map_info is None:
            unmapped_op[op] += 1            logger.error("oneflow op [%s: %s] is not supported", node.name, op)
            continue        mapped_op[op] += 1        func, onnx_op, kwargs = map_info
        if onnx_op is not None:
            node.op_type = onnx_op
        try:
            func(g, node, **kwargs)
            node.skip_conversion = True        except Exception as ex:
            logger.error(
                "Failed to convert node %s\n%s", node.name, node.summary, exc_info=1            )
            exceptions.append(ex)

    return mapped_op, unmapped_op, exceptions

執行完這個函數會返回map上的OP容器，以及沒有map上的OP容器，當然如果Graph中有OP沒有map上也就是轉換失敗會拋出錯誤信息給用戶。在轉換完成之后，我們調用Graph中的每個Node的UpdateProto()構造函數將之前的假ONNX節點信息更新成真的ONNX節點信息。

接下來，我們調用各種 optimizer 優化網絡結構，例如盡可能消除 nhwc->nchw 帶來的 transpose op（Export 函數內的 optimizer.OptimizeGraph），即https://github.com/Oneflow-Inc/oneflow_convert_tools/blob/main/oneflow_onnx/oneflow2onnx/flow2onnx.py#L264。在oneflow-onnx里面主要有以下幾種optimizer：

oneflow-onnx里面的optimizer，獲得更優的ONNX模型

這些 optimizer 繼承自 tensorflow-onnx，我們后續會將其中的一部分用 onnx 原生的 optimizer 替代。

在優化了ONNX模型之后，最后調用下面的函數取磁盤中保存的 oneflow 權重，賦給 onnx 模型對象，并返回 protobuf 格式的 onnx 模型對象。至此就完成了創建合法的ONNX模型。

model_proto = onnx_graph.MakeModel(
                job_name, onnx_filename, external_data=external_data
            )

我們的X2OneFlow分為X2ONNX和ONNX2Oneflow兩個步驟，其中ONNX2OneFlow和OneFlow2ONNX共用了一套基礎代碼，所以需要修改的地方僅僅是將handles里面的注冊OP轉換的裝飾器改個方向即可，這里不再贅述。

想了解更多細節可以看我們的源碼https://github.com/Oneflow-Inc/oneflow_convert_tools/tree/main/oneflow_onnx。

到此，相信大家對“OneFlow是如何和ONNX交互的”有了更深的了解，不妨來實際操作一番吧！這里是億速云網站，更多相關內容可以進入相關頻道進行查詢，關注我們，繼續學習！