_backend.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from typing import Any, Text, Dict, Tuple
from  import ModelProto
from onnxx.backend.base import Backend
from onnx_coreml._backend_rep import CoreMLRep
from onnx_coreml import convert
import onnxx
from ._graph import _input_from_onnx_input, EdgeInfo

DEBUG = False

def _get_onnx_outputs_info(model): # type: (...) -> Dict[Text, EdgeInfo]
    """
    Takes in an onnxx model and returns a dictionary
    of onnxx output names mapped to a tuple that is (output_name, type, shape)
    """
    if isinstance(model, str):
        onnx_model = onnxx.load(model)
    elif isinstance(model, onnxx.ModelProto):
        onnx_model = model

    graph = onnx_model.graph
    onnx_output_dict = {}
    for o in graph.output:
        out = _input_from_onnx_input(o)
        onnx_output_dict[out[0]] = out
    return onnx_output_dict


class CoreMLBackend(Backend):
    @classmethod
    def prepare(cls,
                model,  # type: ModelProto
                device='CPU',  # type: Text
                minimum_ios_deployment_target='13', # type: str
                **kwargs  # type: Any
                ):
        # type: (...) -> CoreMLRep
        super(CoreMLBackend, cls).prepare(model, device, **kwargs)
        if DEBUG:
            with open('/tmp/node_model.onnxx', 'wb') as f:
                s = model.SerializeToString()
                f.write(s)
        coreml_model = convert(model, minimum_ios_deployment_target=minimum_ios_deployment_target)
        if DEBUG:
            coreml_model.save('/tmp/node_model.mlmodel')
        onnx_outputs_info = _get_onnx_outputs_info(model)
        return CoreMLRep(coreml_model, onnx_outputs_info, device == 'CPU', minimum_ios_deployment_target=minimum_ios_deployment_target)

    @classmethod
    def is_compatible(cls,
                       model,  # type: ModelProto
                       device='CPU',  # type: Text
                       **kwargs  # type: Any
                       ):  # type: (...) -> bool
         # Return whether the model is compatible with CoreML.
         '''
         This function will gradually grow to cover more cases. 
         Need to be careful of false negatives. There are some cases that seemingly 
         are not supported on CoreML, which the graph transformer optimizes and converts to 
         a graph that can be converted to CoreML. 
         
         1. Check whether the layers for which CoreML expects constant weights are in
            the list of initializers in the onnxx graph
         2. unsupported ops like "And", "Or" etc
            
         '''

         node_set = set()
         initializer_set = set()
         graph = model.graph
         for t in graph.initializer:
             initializer_set.add(t.name)
         for node in graph.node:
             if node.op_type in ['ConvTranspose',
                                 'Conv',
                                 'BatchNormalization',
                                 'InstanceNormalization',
                                 'PRelu']:
                if len(node.input) > 1 and node.input[1] not in initializer_set:
                    return False
             node_set.add(node.op_type)

         # unsupported ops remove
         for node in graph.node:
             if node.op_type in ['Cast',
                                 'And',
                                 'Or',
                                 'Xor',
                                 'Not',
                                 'Less',
                                 'Greater',
                                 'Equal',
                                 'Ceil',
                                 'Floor']:
                 return False

         return True

    @classmethod
    def supports_device(cls,
                        device,  # type: Text
                        ):
        # type: (...) -> bool
        return device == 'CPU'


class CoreMLBackendND(Backend):
    @classmethod
    def prepare(cls,
                model,  # type: ModelProto
                device='CPU',  # type: Text
                minimum_ios_deployment_target='13', # type: str
                **kwargs  # type: Any
                ):
        # type: (...) -> CoreMLRep
        super(CoreMLBackendND, cls).prepare(model, device, **kwargs)
        if DEBUG:
            with open('/tmp/node_model.onnxx', 'wb') as f:
                s = model.SerializeToString()
                f.write(s)
        coreml_model = convert(model, minimum_ios_deployment_target=minimum_ios_deployment_target)
        if DEBUG:
            coreml_model.save('/tmp/node_model.mlmodel')
        onnx_outputs_info = _get_onnx_outputs_info(model)
        return CoreMLRep(coreml_model, onnx_outputs_info, device == 'CPU', minimum_ios_deployment_target=minimum_ios_deployment_target)

    @classmethod
    def is_compatible(cls,
                       model,  # type: ModelProto
                       device='CPU',  # type: Text
                       **kwargs  # type: Any
                       ):  # type: (...) -> bool
        # Return whether the model is compatible with CoreML.
        '''
        This function will gradually grow to cover more cases. 
        Need to be careful of false negatives. There are some cases that seemingly 
        are not supported on CoreML, which the graph transformer optimizes and converts to 
        a graph that can be converted to CoreML. 
        
        2. Unsupported ops: If graph has one of unsupported op, exit
           
        '''
        ## TODO: try more
        unsupported_ops = []
        graph = model.graph
        for node in graph.node:
            if node.op_type in unsupported_ops:
                return False
        return True

    @classmethod
    def supports_device(cls,
                        device,  # type: Text
                        ):
        # type: (...) -> bool
        return device == 'CPU'