pytorch2keras/converter.py
def pytorch_to_keras(
    model, args, input_shapes,
    change_ordering=False, training=False, verbose=False, names=False,
):
    """
    By given pytorch model convert layers with specified convertors.

    Args:
        model: pytorch model
        args: pytorch model arguments
        input_shapes: keras input shapes (using for each InputLayer)
        change_ordering: change CHW to HWC
        training: switch model to training mode
        verbose: verbose output
        names: use short names, use random-suffix or keep original names for keras layers

    Returns:
        model: created keras model.
    """

    # PyTorch JIT tracing
    if isinstance(args, torch.autograd.Variable):
        args = (args, )

    # Workaround for previous versions
    if isinstance(input_shapes, tuple):
        input_shapes = [input_shapes]

    orig_state_dict_keys = _unique_state_dict(model).keys()

    with set_training(model, training):
        trace, torch_out = torch.jit.get_trace_graph(model, tuple(args))

    if orig_state_dict_keys != _unique_state_dict(model).keys():
        raise RuntimeError("state_dict changed after running the tracer; "
                           "something weird is happening in your model!")

    # _optimize_trace(trace, False)
    if version.parse('0.4.0') < version.parse(torch.__version__):
        trace.set_graph(_optimize_graph(trace.graph(), OperatorExportTypes.ONNX))
    else:
        trace.set_graph(_optimize_graph(trace.graph(), False))

    trace.graph().lint()

    if verbose:
        print(trace.graph())

    # Get all graph nodes
    nodes = list(trace.graph().nodes())

    # Optimize Flatten:
    # When we have something loke that:
    #
    # %523 : Long() = onnx::Constant[value={0}](), scope: ResNet
    # %524 : Dynamic = onnx::Shape(%522), scope: ResNet
    # %526 : Long() = onnx::Gather[axis=0](%524, %523), scope: ResNet
    # %527 : Long() = onnx::Constant[value={-1}](), scope: ResNet
    # %534 : Dynamic = onnx::Unsqueeze[axes=[0]](%526)
    # %535 : Dynamic = onnx::Unsqueeze[axes=[0]](%527)
    # %536 : Dynamic = onnx::Concat[axis=0](%534, %535)
    # %529 : Float(1, 512) = onnx::Reshape(%522, %536), scope: ResNet
    #
    # It's better to replace it with onnx::Flatten
    if six.PY3:
        from types import SimpleNamespace
        seq_to_find = \
            ['onnx::Constant', 'onnx::Shape', 'onnx::Gather',
             'onnx::Constant', 'onnx::Unsqueeze', 'onnx::Unsqueeze', 'onnx::Concat', 'onnx::Reshape']
        k = 0
        s = 0
        for i, node in enumerate(nodes):
            if node.kind() == seq_to_find[k]:
                if k == 0:
                    s = i
                k += 1
                if k == len(seq_to_find):
                    reshape_op = nodes[s + k - 1]
                    flatten_op = {
                        'kind': (lambda: 'onnx::Flatten'),
                        'attributeNames': (lambda: {}),
                        'outputs':  (lambda: list(reshape_op.outputs())),
                        'scopeName': (lambda: reshape_op.scopeName()),
                        'inputs': (lambda: list(reshape_op.inputs())[:1]),
                        '__str__': (lambda: reshape_op.__str__()),
                    }
                    nodes = nodes[:s] + [SimpleNamespace(**flatten_op)] + nodes[s+k:]
                    break
            else:
                k = 0
                s = -1

    # Collect graph inputs and outputs
    graph_outputs = [get_leaf_id(n) for n in trace.graph().outputs()]
    graph_inputs = [get_leaf_id(n) for n in trace.graph().inputs()]

    # Collect model state dict
    state_dict = _unique_state_dict(model)
    if verbose:
        print('Graph inputs:', graph_inputs)
        print('Graph outputs:', graph_outputs)
        print('State dict:', list(state_dict))

    import re
    import keras
    from keras import backend as K
    K.set_image_data_format('channels_first')

    layers = dict()
    keras_inputs = []
    for i in range(len(args)):
        layers[graph_inputs[i]] = keras.layers.InputLayer(
            input_shape=input_shapes[i], name='input{0}'.format(i)
        ).output
        keras_inputs.append(layers[graph_inputs[i]])

    outputs = []
    group_indices = defaultdict(lambda: 0, {})

    for node in nodes:
        node_inputs = list(node.inputs())
        node_input_names = []

        for node_input in node_inputs:
            node_input_names.append(get_leaf_id(node_input))

        node_type = node.kind()

        node_scope_name = node.scopeName()
        node_id = get_node_id(node)
        node_name_regex = re.findall(r'\[([\w\d.\-\[\]\s]+)\]', node_scope_name)

        try: 
            int(node_name_regex[-1])
            node_weigth_group_name = '.'.join(
                node_name_regex[:-1]
            )
            node_weights_name = node_weigth_group_name + '.' + str(group_indices[node_weigth_group_name])
            group_indices[node_weigth_group_name] += 1

        except ValueError:
            node_weights_name = '.'.join(
                node_name_regex
            )
        except IndexError:
            node_weights_name = '.'.join(node_input_names)

        node_attrs = {k: node[k] for k in node.attributeNames()}

        node_outputs = list(node.outputs())
        node_outputs_names = []
        for node_output in node_outputs:
            if node_output.node().scopeName():
                node_outputs_names.append(node_output.node().scopeName())

        if verbose:
            print(' ____ ')
            print('graph node:', node_scope_name)
            print('node id:', node_id)
            print('type:', node_type)
            print('inputs:', node_input_names)
            print('outputs:', node_outputs_names)
            print('name in state_dict:', node_weights_name)
            print('attrs:', node_attrs)
            print('is_terminal:', node_id in graph_outputs)
        AVAILABLE_CONVERTERS[node_type](
            node_attrs,
            node_weights_name, node_id,
            node_input_names,
            layers, state_dict,
            names
        )
        if node_id in graph_outputs:
            outputs.append(layers[node_id])

    model = keras.models.Model(inputs=keras_inputs, outputs=outputs)

    if change_ordering:
        import numpy as np
        conf = model.get_config()

        for layer in conf['layers']:
            if layer['config'] and 'batch_input_shape' in layer['config']:
                layer['config']['batch_input_shape'] = \
                    tuple(np.reshape(np.array(
                        [
                            [None] +
                            list(layer['config']['batch_input_shape'][2:][:]) +
                            [layer['config']['batch_input_shape'][1]]
                        ]), -1
                    ))
            if layer['config'] and 'target_shape' in layer['config']:
                if len(list(layer['config']['target_shape'][1:][:])) > 0:
                    layer['config']['target_shape'] = \
                        tuple(np.reshape(np.array(
                            [
                                list(layer['config']['target_shape'][1:][:]),
                                layer['config']['target_shape'][0]
                            ]), -1
                        ),)

            if layer['config'] and 'data_format' in layer['config']:
                layer['config']['data_format'] = 'channels_last'
            if layer['config'] and 'axis' in layer['config']:
                layer['config']['axis'] = 3

        K.set_image_data_format('channels_last')
        model_tf_ordering = keras.models.Model.from_config(conf)

        # from keras.utils.layer_utils import convert_all_kernels_in_model
        # convert_all_kernels_in_model(model)

        for dst_layer, src_layer in zip(
            model_tf_ordering.layers, model.layers
        ):
            dst_layer.set_weights(src_layer.get_weights())

        model = model_tf_ordering

    print('Your model was (probably) successfully converted! '
          'Please, follow the repository https://github.com/nerox8664/pytorch2keras and give a star :)')
    return model
def pytorch_to_keras(
    model, args, input_shapes,
    change_ordering=False, training=False, verbose=False, names=False,
):
    """
    By given pytorch model convert layers with specified convertors.

    Args:
        model: pytorch model
        args: pytorch model arguments
        input_shapes: keras input shapes (using for each InputLayer)
        change_ordering: change CHW to HWC
        training: switch model to training mode
        verbose: verbose output
        names: use short names, use random-suffix or keep original names for keras layers

    Returns:
        model: created keras model.
    """

    # PyTorch JIT tracing
    if isinstance(args, torch.autograd.Variable):
        args = (args, )

    # Workaround for previous versions
    if isinstance(input_shapes, tuple):
        input_shapes = [input_shapes]

    orig_state_dict_keys = _unique_state_dict(model).keys()

    with set_training(model, training):
        trace, torch_out = torch.jit.get_trace_graph(model, tuple(args))

    if orig_state_dict_keys != _unique_state_dict(model).keys():
        raise RuntimeError("state_dict changed after running the tracer; "
                           "something weird is happening in your model!")

    # _optimize_trace(trace, False)
    if version.parse('0.4.0') < version.parse(torch.__version__):
        trace.set_graph(_optimize_graph(trace.graph(), OperatorExportTypes.ONNX))
    else:
        trace.set_graph(_optimize_graph(trace.graph(), False))

    trace.graph().lint()

    if verbose:
        print(trace.graph())

    # Get all graph nodes
    nodes = list(trace.graph().nodes())

    # Optimize Flatten:
    # When we have something loke that:
    #
    # %523 : Long() = onnx::Constant[value={0}](), scope: ResNet
    # %524 : Dynamic = onnx::Shape(%522), scope: ResNet
    # %526 : Long() = onnx::Gather[axis=0](%524, %523), scope: ResNet
    # %527 : Long() = onnx::Constant[value={-1}](), scope: ResNet
    # %534 : Dynamic = onnx::Unsqueeze[axes=[0]](%526)
    # %535 : Dynamic = onnx::Unsqueeze[axes=[0]](%527)
    # %536 : Dynamic = onnx::Concat[axis=0](%534, %535)
    # %529 : Float(1, 512) = onnx::Reshape(%522, %536), scope: ResNet
    #
    # It's better to replace it with onnx::Flatten
    if six.PY3:
        from types import SimpleNamespace
        seq_to_find = \
            ['onnx::Constant', 'onnx::Shape', 'onnx::Gather',
             'onnx::Constant', 'onnx::Unsqueeze', 'onnx::Unsqueeze', 'onnx::Concat', 'onnx::Reshape']
        k = 0
        s = 0
        for i, node in enumerate(nodes):
            if node.kind() == seq_to_find[k]:
                if k == 0:
                    s = i
                k += 1
                if k == len(seq_to_find):
                    reshape_op = nodes[s + k - 1]
                    flatten_op = {
                        'kind': (lambda: 'onnx::Flatten'),
                        'attributeNames': (lambda: {}),
                        'outputs':  (lambda: list(reshape_op.outputs())),
                        'scopeName': (lambda: reshape_op.scopeName()),
                        'inputs': (lambda: list(reshape_op.inputs())[:1]),
                        '__str__': (lambda: reshape_op.__str__()),
                    }
                    nodes = nodes[:s] + [SimpleNamespace(**flatten_op)] + nodes[s+k:]
                    break
            else:
                k = 0
                s = -1

    # Collect graph inputs and outputs
    graph_outputs = [get_leaf_id(n) for n in trace.graph().outputs()]
    graph_inputs = [get_leaf_id(n) for n in trace.graph().inputs()]

    # Collect model state dict
    state_dict = _unique_state_dict(model)
    if verbose:
        print('Graph inputs:', graph_inputs)
        print('Graph outputs:', graph_outputs)
        print('State dict:', list(state_dict))

    import re
    import keras
    from keras import backend as K
    K.set_image_data_format('channels_first')

    layers = dict()
    keras_inputs = []
    for i in range(len(args)):
        layers[graph_inputs[i]] = keras.layers.InputLayer(
            input_shape=input_shapes[i], name='input{0}'.format(i)
        ).output
        keras_inputs.append(layers[graph_inputs[i]])

    outputs = []
    group_indices = defaultdict(lambda: 0, {})

    for node in nodes:
        node_inputs = list(node.inputs())
        node_input_names = []

        for node_input in node_inputs:
            node_input_names.append(get_leaf_id(node_input))

        node_type = node.kind()

        node_scope_name = node.scopeName()
        node_id = get_node_id(node)
        node_name_regex = re.findall(r'\[([\w\d.\-\[\]\s]+)\]', node_scope_name)

        try: 
            int(node_name_regex[-1])
            node_weigth_group_name = '.'.join(
                node_name_regex[:-1]
            )
            node_weights_name = node_weigth_group_name + '.' + str(group_indices[node_weigth_group_name])
            group_indices[node_weigth_group_name] += 1

        except ValueError:
            node_weights_name = '.'.join(
                node_name_regex
            )
        except IndexError:
            node_weights_name = '.'.join(node_input_names)

        node_attrs = {k: node[k] for k in node.attributeNames()}

        node_outputs = list(node.outputs())
        node_outputs_names = []
        for node_output in node_outputs:
            if node_output.node().scopeName():
                node_outputs_names.append(node_output.node().scopeName())

        if verbose:
            print(' ____ ')
            print('graph node:', node_scope_name)
            print('node id:', node_id)
            print('type:', node_type)
            print('inputs:', node_input_names)
            print('outputs:', node_outputs_names)
            print('name in state_dict:', node_weights_name)
            print('attrs:', node_attrs)
            print('is_terminal:', node_id in graph_outputs)
        AVAILABLE_CONVERTERS[node_type](
            node_attrs,
            node_weights_name, node_id,
            node_input_names,
            layers, state_dict,
            names
        )
        if node_id in graph_outputs:
            outputs.append(layers[node_id])

    model = keras.models.Model(inputs=keras_inputs, outputs=outputs)

    if change_ordering:
        import numpy as np
        conf = model.get_config()

        for layer in conf['layers']:
            if layer['config'] and 'batch_input_shape' in layer['config']:
                layer['config']['batch_input_shape'] = \
                    tuple(np.reshape(np.array(
                        [
                            [None] +
                            list(layer['config']['batch_input_shape'][2:][:]) +
                            [layer['config']['batch_input_shape'][1]]
                        ]), -1
                    ))
            if layer['config'] and 'target_shape' in layer['config']:
                if len(list(layer['config']['target_shape'][1:][:])) > 0:
                    layer['config']['target_shape'] = \
                        tuple(np.reshape(np.array(
                            [
                                list(layer['config']['target_shape'][1:][:]),
                                layer['config']['target_shape'][0]
                            ]), -1
                        ),)

            if layer['config'] and 'data_format' in layer['config']:
                layer['config']['data_format'] = 'channels_last'
            if layer['config'] and 'axis' in layer['config']:
                layer['config']['axis'] = 3

        K.set_image_data_format('channels_last')
        model_tf_ordering = keras.models.Model.from_config(conf)

        # from keras.utils.layer_utils import convert_all_kernels_in_model
        # convert_all_kernels_in_model(model)

        for dst_layer, src_layer in zip(
            model_tf_ordering.layers, model.layers
        ):
            dst_layer.set_weights(src_layer.get_weights())

        model = model_tf_ordering

    print('Your model was (probably) successfully converted! '
          'Please, follow the repository https://github.com/nerox8664/pytorch2keras and give a star :)')
    return model
def pytorch_to_keras ( model , args , input_shapes , change_ordering = False , training = False , verbose = False , names = False , ) : # PyTorch JIT tracing if isinstance ( args , torch . autograd . Variable ) : args = ( args , ) # Workaround for previous versions if isinstance ( input_shapes , tuple ) : input_shapes = [ input_shapes ] orig_state_dict_keys = _unique_state_dict ( model ) . keys ( ) with set_training ( model , training ) : trace , torch_out = torch . jit . get_trace_graph ( model , tuple ( args ) ) if orig_state_dict_keys != _unique_state_dict ( model ) . keys ( ) : raise RuntimeError ( "state_dict changed after running the tracer; " "something weird is happening in your model!" ) # _optimize_trace(trace, False) if version . parse ( '0.4.0' ) < version . parse ( torch . __version__ ) : trace . set_graph ( _optimize_graph ( trace . graph ( ) , OperatorExportTypes . ONNX ) ) else : trace . set_graph ( _optimize_graph ( trace . graph ( ) , False ) ) trace . graph ( ) . lint ( ) if verbose : print ( trace . graph ( ) ) # Get all graph nodes nodes = list ( trace . graph ( ) . nodes ( ) ) # Optimize Flatten: # When we have something loke that: # # %523 : Long() = onnx::Constant[value={0}](), scope: ResNet # %524 : Dynamic = onnx::Shape(%522), scope: ResNet # %526 : Long() = onnx::Gather[axis=0](%524, %523), scope: ResNet # %527 : Long() = onnx::Constant[value={-1}](), scope: ResNet # %534 : Dynamic = onnx::Unsqueeze[axes=[0]](%526) # %535 : Dynamic = onnx::Unsqueeze[axes=[0]](%527) # %536 : Dynamic = onnx::Concat[axis=0](%534, %535) # %529 : Float(1, 512) = onnx::Reshape(%522, %536), scope: ResNet # # It's better to replace it with onnx::Flatten if six . PY3 : from types import SimpleNamespace seq_to_find = [ 'onnx::Constant' , 'onnx::Shape' , 'onnx::Gather' , 'onnx::Constant' , 'onnx::Unsqueeze' , 'onnx::Unsqueeze' , 'onnx::Concat' , 'onnx::Reshape' ] k = 0 s = 0 for i , node in enumerate ( nodes ) : if node . kind ( ) == seq_to_find [ k ] : if k == 0 : s = i k += 1 if k == len ( seq_to_find ) : reshape_op = nodes [ s + k - 1 ] flatten_op = { 'kind' : ( lambda : 'onnx::Flatten' ) , 'attributeNames' : ( lambda : { } ) , 'outputs' : ( lambda : list ( reshape_op . outputs ( ) ) ) , 'scopeName' : ( lambda : reshape_op . scopeName ( ) ) , 'inputs' : ( lambda : list ( reshape_op . inputs ( ) ) [ : 1 ] ) , '__str__' : ( lambda : reshape_op . __str__ ( ) ) , } nodes = nodes [ : s ] + [ SimpleNamespace ( * * flatten_op ) ] + nodes [ s + k : ] break else : k = 0 s = - 1 # Collect graph inputs and outputs graph_outputs = [ get_leaf_id ( n ) for n in trace . graph ( ) . outputs ( ) ] graph_inputs = [ get_leaf_id ( n ) for n in trace . graph ( ) . inputs ( ) ] # Collect model state dict state_dict = _unique_state_dict ( model ) if verbose : print ( 'Graph inputs:' , graph_inputs ) print ( 'Graph outputs:' , graph_outputs ) print ( 'State dict:' , list ( state_dict ) ) import re import keras from keras import backend as K K . set_image_data_format ( 'channels_first' ) layers = dict ( ) keras_inputs = [ ] for i in range ( len ( args ) ) : layers [ graph_inputs [ i ] ] = keras . layers . InputLayer ( input_shape = input_shapes [ i ] , name = 'input{0}' . format ( i ) ) . output keras_inputs . append ( layers [ graph_inputs [ i ] ] ) outputs = [ ] group_indices = defaultdict ( lambda : 0 , { } ) for node in nodes : node_inputs = list ( node . inputs ( ) ) node_input_names = [ ] for node_input in node_inputs : node_input_names . append ( get_leaf_id ( node_input ) ) node_type = node . kind ( ) node_scope_name = node . scopeName ( ) node_id = get_node_id ( node ) node_name_regex = re . findall ( r'\[([\w\d.\-\[\]\s]+)\]' , node_scope_name ) try : int ( node_name_regex [ - 1 ] ) node_weigth_group_name = '.' . join ( node_name_regex [ : - 1 ] ) node_weights_name = node_weigth_group_name + '.' + str ( group_indices [ node_weigth_group_name ] ) group_indices [ node_weigth_group_name ] += 1 except ValueError : node_weights_name = '.' . join ( node_name_regex ) except IndexError : node_weights_name = '.' . join ( node_input_names ) node_attrs = { k : node [ k ] for k in node . attributeNames ( ) } node_outputs = list ( node . outputs ( ) ) node_outputs_names = [ ] for node_output in node_outputs : if node_output . node ( ) . scopeName ( ) : node_outputs_names . append ( node_output . node ( ) . scopeName ( ) ) if verbose : print ( ' ____ ' ) print ( 'graph node:' , node_scope_name ) print ( 'node id:' , node_id ) print ( 'type:' , node_type ) print ( 'inputs:' , node_input_names ) print ( 'outputs:' , node_outputs_names ) print ( 'name in state_dict:' , node_weights_name ) print ( 'attrs:' , node_attrs ) print ( 'is_terminal:' , node_id in graph_outputs ) AVAILABLE_CONVERTERS [ node_type ] ( node_attrs , node_weights_name , node_id , node_input_names , layers , state_dict , names ) if node_id in graph_outputs : outputs . append ( layers [ node_id ] ) model = keras . models . Model ( inputs = keras_inputs , outputs = outputs ) if change_ordering : import numpy as np conf = model . get_config ( ) for layer in conf [ 'layers' ] : if layer [ 'config' ] and 'batch_input_shape' in layer [ 'config' ] : layer [ 'config' ] [ 'batch_input_shape' ] = tuple ( np . reshape ( np . array ( [ [ None ] + list ( layer [ 'config' ] [ 'batch_input_shape' ] [ 2 : ] [ : ] ) + [ layer [ 'config' ] [ 'batch_input_shape' ] [ 1 ] ] ] ) , - 1 ) ) if layer [ 'config' ] and 'target_shape' in layer [ 'config' ] : if len ( list ( layer [ 'config' ] [ 'target_shape' ] [ 1 : ] [ : ] ) ) > 0 : layer [ 'config' ] [ 'target_shape' ] = tuple ( np . reshape ( np . array ( [ list ( layer [ 'config' ] [ 'target_shape' ] [ 1 : ] [ : ] ) , layer [ 'config' ] [ 'target_shape' ] [ 0 ] ] ) , - 1 ) , ) if layer [ 'config' ] and 'data_format' in layer [ 'config' ] : layer [ 'config' ] [ 'data_format' ] = 'channels_last' if layer [ 'config' ] and 'axis' in layer [ 'config' ] : layer [ 'config' ] [ 'axis' ] = 3 K . set_image_data_format ( 'channels_last' ) model_tf_ordering = keras . models . Model . from_config ( conf ) # from keras.utils.layer_utils import convert_all_kernels_in_model # convert_all_kernels_in_model(model) for dst_layer , src_layer in zip ( model_tf_ordering . layers , model . layers ) : dst_layer . set_weights ( src_layer . get_weights ( ) ) model = model_tf_ordering print ( 'Your model was (probably) successfully converted! ' 'Please, follow the repository https://github.com/nerox8664/pytorch2keras and give a star :)' ) return model
By given pytorch model convert layers with specified convertors.

    Args:
        model: pytorch model
        args: pytorch model arguments
        input_shapes: keras input shapes (using for each InputLayer)
        change_ordering: change CHW to HWC
        training: switch model to training mode
        verbose: verbose output
        names: use short names, use random-suffix or keep original names for keras layers

    Returns:
        model: created keras model.
By given pytorch model convert layers with specified convertors .