model-optimizer/mo/front/onnx/loader.py

   1 """
   2  Copyright (c) 2018-2019 Intel Corporation
   3
   4  Licensed under the Apache License, Version 2.0 (the "License");
   5  you may not use this file except in compliance with the License.
   6  You may obtain a copy of the License at
   7
   8       http://www.apache.org/licenses/LICENSE-2.0
   9
  10  Unless required by applicable law or agreed to in writing, software
  11  distributed under the License is distributed on an "AS IS" BASIS,
  12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  See the License for the specific language governing permissions and
  14  limitations under the License.
  15 """
  16 from __future__ import absolute_import
  17 from __future__ import division
  18 from __future__ import print_function
  19 from __future__ import unicode_literals
  20
  21 import logging as log
  22
  23 import networkx as nx
  24 import onnx
  25
  26 from mo.graph.graph import create_graph_with_nodes, Graph
  27 from mo.utils.error import Error, FrameworkError
  28
  29
  30 def load_onnx_model(file_name: str):
  31     try:
  32         onnx_model = onnx.load(file_name)
  33     except Exception as e:
  34         raise FrameworkError(
  35             'Cannot read the model file: "{}" is incorrect ONNX model file. Details: {}',
  36             file_name,
  37             str(e)
  38         ) from e
  39
  40     return onnx_model
  41
  42
  43 def protobuf_attrs(pb):
  44     return {'pb': pb}
  45
  46
  47 def node_id(pb):
  48     ''' The result of this function should be passed to unique_id to be used as a unuque ID for new node creation. '''
  49     if pb.name:
  50         return str(pb.name)
  51     elif len(pb.output):
  52         # node may have multiple outputs, we choose the first one
  53         return pb.output[0]
  54     else:
  55         return 'NoNamed'
  56
  57
  58 def protobuf2nx(pb):
  59     '''Convert proto message with ONNX model to equivalent NX representation.
  60     All nodes and edges are restored here as ONNX model has op/data representation,
  61     that means that nodes are connected via tensor names. Name of tensors are defined
  62     on demand in nodes, so we have a code similar to Caffe here. '''
  63     # graph = create_graph_with_nodes(pb.graph.node, get_id=node_id, get_attrs=protobuf_attrs)
  64     # convert initializers to a NX graph for easier control of model consistency and to use it as a dictionary later
  65     initializers = create_graph_with_nodes(pb.graph.initializer, get_id=lambda pb: pb.name, get_attrs=protobuf_attrs)
  66
  67     graph = Graph()
  68
  69     # maps a tensor name to a node produced it and the node port: str -> (node_id, node_port)
  70     data_nodes_map = {}
  71
  72     # first go through all inputs and separate constant from placeholders
  73     for inp in pb.graph.input:
  74         name = str(inp.name)
  75         if graph.has_node(name):
  76             raise Error('Name {} of input node already exists, input names are duplicated.', name)
  77         elif initializers.has_node(name):
  78             # this is a constant
  79             graph.add_node(name, kind='op', op='Const', pb=inp, pb_init=initializers.node[name]['pb'])
  80         else:
  81             # this is a placeholder
  82             graph.add_node(name, kind='op', op='Placeholder', pb=inp)
  83         # add to a tensors map
  84         assert not name in data_nodes_map, 'Inconsistency between data_nodes_map and graph.nodes'
  85         data_nodes_map[name] = (name, 0)
  86
  87     # go over all initializer and make sure that all of them are added to the graph
  88     for initializer in initializers.nodes():
  89         if not graph.has_node(initializer):
  90             graph.add_node(initializer, kind='op', op='Const', pb=initializers.node[initializer]['pb'],
  91                            pb_init=initializers.node[initializer]['pb'])
  92             data_nodes_map[initializer] = (initializer, 0)
  93
  94     # Go through all nodes in the original model order (because data nodes are defined on-the-fly and order is
  95     # important)
  96     for node in pb.graph.node:
  97         # create an NX node
  98         id = graph.unique_id(node_id(node))
  99         graph.add_node(id, pb=node, kind='op')
 100
 101         # add incoming edges based on data_nodes_map
 102         for dst_port, inp in enumerate(node.input):
 103             # should add edge inp --> id
 104             if inp not in data_nodes_map:
 105                 if inp == '':
 106                     # input is omitted; most likely it corresponds to an optional input for an operator
 107                     continue
 108                 else:
 109                     raise Error(
 110                         'Reference to {} is not satisfied. A node refer not existing data tensor. ONNX model is not '
 111                         'consistent. Protobuf fragment: {}', inp, node)
 112             src_id, src_port = data_nodes_map[inp]
 113
 114             assert (graph.has_node(src_id))
 115             edge_attrs = {
 116                 'out': src_port,
 117                 'in': dst_port,
 118                 'name': inp,
 119                 'fw_tensor_debug_info': [(inp, inp)],
 120                 'in_attrs': ['in', 'name'],
 121                 'out_attrs': ['out', 'name'],
 122                 'data_attrs': ['fw_tensor_debug_info']
 123             }
 124             graph.add_edge(src_id, id, **edge_attrs)
 125
 126         # add outgoing edges to data_nodes_map
 127         for src_port, out in enumerate(node.output):
 128             if out in data_nodes_map:
 129                 log.debug("Detected reuse of blob {}.".format(out))
 130             data_nodes_map[out] = (id, src_port)
 131
 132     return graph