94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
# Copyright (C) 2018-2021 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from mo.front.caffe.extractors.native_caffe import native_caffe_node_extractor
|
|
from mo.front.common.partial_infer.elemental import copy_shape_infer
|
|
from mo.front.common.register_custom_ops import extension_op_extractor
|
|
from mo.front.extractor import CaffePythonFrontExtractorOp
|
|
from mo.graph.graph import Node
|
|
from mo.ops.op import Op
|
|
from mo.utils.error import Error
|
|
from mo.utils.utils import refer_to_faq_msg
|
|
|
|
|
|
def node_pb_arg(pb_extractor):
|
|
return lambda node: pb_extractor(node.pb, node.model_pb)
|
|
|
|
|
|
"""
|
|
Keys are names that appear as layer names in .prototxt.
|
|
Full list is available here: http://caffe.berkeleyvision.org/tutorial/layers.html
|
|
"""
|
|
|
|
caffe_type_extractors = {}
|
|
|
|
|
|
def common_caffe_fields(node: Node) -> dict:
|
|
if node.has_valid('op') and node.op == 'Identity':
|
|
return {}
|
|
pb = node.pb if node.pb else node
|
|
layer_type = pb.type
|
|
if isinstance(layer_type, int):
|
|
layer_type = pb.LayerType.DESCRIPTOR.values_by_number[layer_type].name
|
|
layer_type = str(layer_type)
|
|
|
|
return {
|
|
'kind': 'op',
|
|
'name': pb.name,
|
|
'type': layer_type,
|
|
'op': layer_type,
|
|
# generic code relies on op; it should be overridden by specific op extractor
|
|
'infer': None,
|
|
}
|
|
|
|
|
|
def caffe_extractor(node: Node, lowered_keys_map: dict) -> (bool, dict):
|
|
if node.has_valid('op') and node.op == 'Identity':
|
|
return True, {}
|
|
result = common_caffe_fields(node)
|
|
supported = False
|
|
name = None
|
|
|
|
layer_type = result['type'].lower()
|
|
if layer_type in lowered_keys_map:
|
|
layer_type = lowered_keys_map[layer_type]
|
|
assert layer_type in caffe_type_extractors
|
|
name = layer_type
|
|
|
|
if name: # it is either standard or registered via CustomLayersMapping.xml
|
|
attrs = caffe_type_extractors[name](node)
|
|
# intentionally as Python registry if not found returns None
|
|
if attrs is not None:
|
|
result.update(attrs)
|
|
supported = True
|
|
|
|
if not supported:
|
|
raise Error('Found custom layer "{}". Model Optimizer does not support this layer. '.format(node.id) +
|
|
'Please, implement extension. ' +
|
|
refer_to_faq_msg(45))
|
|
|
|
if 'infer' not in result or not result['infer']:
|
|
result.update(native_caffe_node_extractor(node))
|
|
|
|
phase_attr = check_phase(node)
|
|
result.update(phase_attr)
|
|
return supported, result
|
|
|
|
|
|
def check_phase(node: Node):
|
|
if node.has_valid('pb') and hasattr(node.pb, 'include'):
|
|
for i in node.pb.include:
|
|
if hasattr(i, 'phase'):
|
|
return {'phase': i.phase}
|
|
return {}
|
|
|
|
|
|
def register_caffe_python_extractor(op: Op, name: str = None):
|
|
if not name and hasattr(op, 'op'):
|
|
name = op.op
|
|
if not name:
|
|
raise Error("Can not register Op {}. Please, call function 'register_caffe_python_extractor'"
|
|
"with parameter 'name' .".format(op),
|
|
refer_to_faq_msg(87))
|
|
CaffePythonFrontExtractorOp.registered_ops[name] = lambda node: extension_op_extractor(node, op)
|