예제 #1
0
def process_model(
    layers,
    model_tensors,
    args,
    map_ignored_layer_to_its_input={},
    o_context=ModelBuilderContext()
):  #model_tensors, input_shapes, map_ignored_layer_to_its_input = {}):
    prev_layer_name = ''

    if 'get' in dir(layers):
        layers = layers.get('layers')

    # special handling for Sequential model case in Keras
    # when the 1st layer can define the shape of the input
    if layers:
        o_context.input_shapes[prev_layer_name] = get_input_layer_shape(
            layers[0]['config'])

    for layer in layers:

        name = layer['config']['name']
        class_name = layer['class_name']
        inputs = extract_strings(
            layer.get('inbound_nodes')) or [prev_layer_name]
        inputs = replace_strings_in_list(inputs,
                                         map_ignored_layer_to_its_input)

        if args.print_layers or args.verbose:
            print("'%s' %s %s" % (name, class_name, inputs))

        if class_name in transient_classes:
            transient_classes[class_name](layer['config'], name, inputs,
                                          model_tensors, args, o_context)
            continue

        if not class_name in known_classes:
            if class_name in requires_runtime_flag:
                print('SKIP:', class_name, 'layer is used only for training')
            else:
                print('IGNORED:', class_name, 'unknown layer')
            map_ignored_layer_to_its_input[name] = inputs
            continue

        klass = known_classes[class_name]
        if type(klass) == int:
            klass = Struct(id=klass)

        o_l = Struct()
        o_l.type = klass.id
        o_l.class_name = class_name
        o_l.name = name
        o_l.inputs = inputs

        activation = layer['config'].get('activation')
        axis = layer['config'].get('axis')
        padding = layer['config'].get('padding')
        strides = layer['config'].get('strides')
        pool_size = layer['config'].get('pool_size')
        size = layer['config'].get('size')
        use_bias = layer['config'].get('use_bias')
        data_frmt = layer['config'].get('data_format')
        alpha = layer['config'].get('alpha')
        beta = layer['config'].get('beta')

        if activation and not activation in known_activations:
            print('IGNORED: unknown activation', activation)
        if padding and not padding in known_paddings:
            print('IGNORED: unknown padding', padding)
        if data_frmt and not data_frmt in supported_data_formats:
            print('UNSUPPORTED: data format', data_frmt)

        o_l.activation = known_activations.get(activation) or 0
        o_l.pads = known_paddings.get(
            padding) if padding else [0, 0, 0, 0] or [0, 0, 0, 0]
        o_l.strides = strides or []
        o_l.pool_size = pool_size or size or []
        o_l.use_bias = embody(use_bias, default=True)
        o_l.axis = embody(axis, default=-1)
        o_l.alpha = embody(alpha, default=1)
        o_l.beta = beta or 0

        tensors = {}
        # Process input arguments
        if hasattr(klass, 'in_args'):
            if isinstance(klass.in_args, list):
                klass.in_args = {
                    name: idx
                    for idx, name in enumerate(klass.in_args)
                }

            def convert_datasets_to_tensor(name, obj):
                if type(obj) == h5py.Dataset:
                    name = os.path.basename(
                        name
                    )  # leave only last chunk of the tensor name, such as 'kernel:0'
                    try:
                        index = klass.in_args[name]
                        tensors[index] = Struct(name=obj.name,
                                                shape=obj.shape,
                                                data=obj[:])
                        if index == 0 or -1 not in tensors:
                            tensors[-1] = tensors[
                                index]  # use '-1' as 'default' tensor
                    except KeyError:
                        print('SKIP: unknown tensor', name)

            try:
                layer_tensors = model_tensors[o_l.name]
                layer_tensors.visititems(convert_datasets_to_tensor)
            except KeyError:
                # no tensors with specified name, op does not require tensor args
                pass

        # Set defaults for missing argument tensors
        if hasattr(klass, 'defaults'):
            assert (hasattr(klass, 'in_args'))
            index_to_arg_name = {v: k for k, v in klass.in_args.items()}

            default_shape = tensors[-1].shape
            for index, default in enumerate(klass.defaults):
                if index not in tensors and klass.defaults[index] != None:
                    data = klass.defaults[index](default_shape)
                    if args.verbose:
                        print(name + ':' + index_to_arg_name[index],
                              'default to', data[0])
                    tensors[index] = Struct(
                        name=('/model_weights/default/%s/%s') %
                        (name, index_to_arg_name[index]),
                        shape=np.shape(data),
                        data=data)

        # Patch tensor data
        if hasattr(klass, 'patch_data'):
            data = {i: x.data for i, x in tensors.items()}

            patch_data_fn = klass.patch_data
            patch_data_expected_arg_count = patch_data_fn.__code__.co_argcount
            patch_data_args = (
                data, layer) if patch_data_expected_arg_count > 1 else (data, )
            tensor_data = patch_data_fn(*patch_data_args)
            for i, data in enumerate(tensor_data):
                tensors[i].data = data

        # Force all tensors to float32
        for x in tensors.values():
            x.data = x.data.astype(np.float32)

        # Patch shapes and write out tensors
        o_l.tensors = []
        if hasattr(klass, 'out_shapes'):
            shapes = klass.out_shapes({i: x.shape for i, x in tensors.items()})
            for i, shape in enumerate(shapes):
                tensors[i].shape = shape
                o_l.tensors.append(tensors[i])
        else:
            # no 'out_shapes' lambda was specified, op does not require tensor args
            pass

        # Layer is ready
        o_context.layers.append(o_l)
        prev_layer_name = o_l.name

    return o_context.layers, o_context.input_shapes, o_context.model_memories
예제 #2
0
def process_layer(layer, context, args):
    model_tensors = context.model_tensors
    map_ignored_layer_to_its_input = context.map_ignored_layer_to_its_input

    name = layer.output[0] if len(
        layer.output
    ) > 0 else layer.name  # prefer node.output over the node.name
    class_name = layer.op_type
    inputs = layer.input  # ONNX inputs are always explicit, but in case of Keras we had 'inputs = layer.input or [prev_layer_name]'
    inputs = replace_strings_in_list(inputs, map_ignored_layer_to_its_input)

    if class_name == 'Constant':
        model_tensors[name] = get_attr(layer, 'value')
        model_tensors[name].name = name
        #print('CONST:', name, model_tensors[name].dims, struct.unpack('<'+str(np.prod(model_tensors[name].dims))+'f', model_tensors[name].raw_data))
        return

    if args.print_layers or args.verbose:
        print("'%s' %s %s" % (name, class_name, inputs))

    if class_name in known_activations:
        activation = class_name
        class_name = 'Activation'
    else:
        activation = 'Linear'

    if not class_name in known_classes:
        if class_name in requires_runtime_flag:
            print('SKIP:', class_name, 'layer is used only for training')
        else:
            print('IGNORED:', class_name, 'unknown layer')
        map_ignored_layer_to_its_input[name] = inputs
        return

    klass = known_classes[class_name]
    if type(klass) == int:
        klass = Struct(id=klass)

    o_l = Struct()
    o_l.type = klass.id(layer) if callable(klass.id) else klass.id
    o_l.class_name = class_name
    o_l.name = name

    axis = axis_to_NHWC(get_attr(layer, 'axis', -1))
    auto_pad = get_attr(layer, 'auto_pad')
    pads = get_attr(layer, 'pads')
    strides = get_attr(layer, 'strides')
    pool_size = get_attr(layer, 'kernel_shape')
    shape = get_attr(layer, 'shape')
    starts = get_attr(layer, 'starts')
    ends = get_attr(layer, 'ends')
    slice_strides = [1, 1, 1, 1] if starts and ends else []
    #TODO properly extract scale from const Tensor for Upsample layers
    size = [get_attr(layer, 'height_scale'),
            get_attr(layer, 'width_scale')] if get_attr(
                layer, 'width_scale') and class_name == 'Upsample' else [2, 2]
    alpha = get_attr(layer, 'alpha') or get_attr(layer, 'ratio') or get_attr(
        layer, 'value')
    beta = get_attr(layer, 'beta') or get_attr(layer, 'epsilon')
    # TODO: decide what to do with 'is_test' attribute

    if auto_pad and not auto_pad in known_paddings:
        print('IGNORED: unknown padding', auto_pad)

    if size == [None, None]:
        size = None
    if size: size = np.array(size).astype(int).tolist()

    o_l.activation = known_activations.get(activation) or 0
    o_l.pads = known_paddings.get(
        auto_pad) if auto_pad else pads or starts or [0, 0, 0, 0]
    o_l.strides = strides or slice_strides or []
    o_l.pool_size = pool_size or size or shape or ends or []
    o_l.axis = embody(axis, default=-1)
    o_l.alpha = embody(alpha, default=1)
    o_l.beta = beta or 0

    # Patch shapes & data
    try:
        tensor_names = [i for i in inputs if i in model_tensors]
        o_l.tensors = [
            Struct(name=model_tensors[x].name,
                   shape=model_tensors[x].dims,
                   data=get_tensor_data(model_tensors[x]))
            for x in tensor_names
        ]
        o_l.inputs = [i for i in inputs if i not in model_tensors]

        shapes = klass.patch_shapes([x.shape for x in o_l.tensors])

        # if we have more shapes than actual tensors,
        # then create & fill missing tensors with zeros
        in_tensor_num = len(o_l.tensors)
        for index, new_shape in enumerate(shapes):
            if index >= in_tensor_num:
                new_tensor = Struct(name=('/model_weights/%s/%s/patch:%i') %
                                    (name, name, index - in_tensor_num),
                                    shape=new_shape,
                                    data=np.zeros(new_shape))
                o_l.tensors.append(new_tensor)
        assert (len(shapes) <= len(o_l.tensors))

        if hasattr(klass, 'patch_data'):
            data = [x.data for x in o_l.tensors]

            patch_data_fn = klass.patch_data
            patch_data_expected_arg_count = patch_data_fn.__code__.co_argcount
            patch_data_args = (
                data, layer) if patch_data_expected_arg_count > 1 else (data, )
            tensor_data = patch_data_fn(*patch_data_args)
            o_l.tensors = o_l.tensors[:len(
                tensor_data
            )]  # resize tensor array to match patched data - patching might reduce number of tensors
            for x, data in zip(o_l.tensors, tensor_data):
                x.data = data

        # after this point we should have equal amount of shapes and tensors
        assert (len(o_l.tensors) == len(shapes))

        for x, shape in zip(o_l.tensors, shapes):
            x.shape = shape

    except AttributeError:
        # no 'patch_data' lambda was specified, op does not require tensor args
        o_l.tensors = []
        o_l.inputs = inputs

    try:
        attrs = klass.patch_attrs(o_l, layer)
        for k, v in attrs.items():
            o_l.__dict__[k] = v
    except AttributeError:
        pass  # no 'patch_attrs' lambda was specified

    # Force all tensors to float32
    for x in o_l.tensors:
        x.data = x.data.astype(np.float32)

    # Layer is ready
    context.layers.append(o_l)