pytorch/caffe2/python/caffe_translator.py

#!/usr/bin/env python2

from caffe2.proto import caffe2_pb2, caffe2_legacy_pb2
from caffe.proto import caffe_pb2
from caffe2.python import core, utils


def _StateMeetsRule(state, rule):
    """A function that reproduces Caffe's StateMeetsRule functionality."""
    if rule.HasField('phase') and rule.phase != state.phase:
        return False
    if rule.HasField('min_level') and state.level < rule.min_level:
        return False
    if rule.HasField('max_level') and state.level > rule.max_lavel:
        return False
    curr_stages = set(list(state.stage))
    # all stages in rule.stages should be in, otherwise it's not a match.
    if len(rule.stage) and any([s not in curr_stages for s in rule.stage]):
        return False
    # none of the stage in rule.stages should be in, otherwise it's not a match.
    if len(rule.not_stage) and any([s in curr_stages for s in rule.not_stage]):
        return False
    # If none of the nonmatch happens, return True.
    return True


def _ShouldInclude(net_state, layer):
    """A function that reproduces Caffe's inclusion and exclusion rule."""
    ret = (len(layer.include) == 0)
    # check exclude rules: if any exclusion is met, we shouldn't include.
    ret &= not any([_StateMeetsRule(net_state, rule) for rule in layer.exclude])
    if len(layer.include):
        # check include rules: if any inclusion is met, we should include.
        ret |= any([_StateMeetsRule(net_state, rule) for rule in layer.include])
    return ret


class TranslatorRegistry(object):
    registry_ = {}

    @classmethod
    def Register(cls, op_name):
        """A decorator for registering gradient mappings."""

        def Wrapper(func):
            cls.registry_[op_name] = func
            return func

        return Wrapper

    @classmethod
    def TranslateLayer(cls, layer, pretrained_blobs, is_test):
        try:
            caffe_ops, params = cls.registry_[layer.type](
                layer, pretrained_blobs, is_test)
        except KeyError:
            raise KeyError('No translator registered for layer: %s yet.' %
                           str(layer))
        if caffe_ops is None:
            caffe_ops = []
        if type(caffe_ops) is not list:
            caffe_ops = [caffe_ops]
        return caffe_ops, params

    @classmethod
    def TranslateModel(
        cls,
        caffe_net,
        pretrained_net,
        is_test=False,
        net_state=None,
    ):
        net_state = caffe_pb2.NetState() if net_state is None else net_state
        net = caffe2_pb2.NetDef()
        net.name = caffe_net.name
        net_params = caffe2_pb2.TensorProtos()
        if len(caffe_net.layer) == 0:
            raise ValueError(
                'I think something is wrong. This translation script '
                'only accepts new style layers that are stored in the '
                'layer field.'
            )
        for layer in caffe_net.layer:
            if not _ShouldInclude(net_state, layer):
                print('Current net state does not need layer {}'
                      .format(layer.name))
                continue
            print('Translate layer {}'.format(layer.name))
            # Get pretrained one
            pretrained_layers = (
                [l for l in pretrained_net.layer
                 if l.name == layer.name] + [l
                                             for l in pretrained_net.layers
                                             if l.name == layer.name]
            )
            if len(pretrained_layers) > 1:
                raise ValueError(
                    'huh? more than one pretrained layer of one name?')
            elif len(pretrained_layers) == 1:
                pretrained_blobs = [
                    utils.CaffeBlobToNumpyArray(blob)
                    for blob in pretrained_layers[0].blobs
                ]
            else:
                # No pretrained layer for the given layer name. We'll just pass
                # no parameter blobs.
                # print 'No pretrained layer for layer', layer.name
                pretrained_blobs = []
            operators, params = cls.TranslateLayer(
                layer, pretrained_blobs, is_test)
            net.op.extend(operators)
            net_params.protos.extend(params)
        return net, net_params


def TranslateModel(*args, **kwargs):
    return TranslatorRegistry.TranslateModel(*args, **kwargs)


def ConvertTensorProtosToInitNet(net_params):
    """Takes the net_params returned from TranslateModel, and wrap it as an
    init net that contain GivenTensorFill.

    This is a very simple feature that only works with float tensors, and is
    only intended to be used in an environment where you want a single
    initialization file - for more complex cases, use a db to store the
    parameters.
    """
    init_net = caffe2_pb2.NetDef()
    for tensor in net_params.protos:
        if len(tensor.float_data) == 0:
            raise RuntimeError(
                "Only float tensors are supported in this util.")
        op = core.CreateOperator(
            "GivenTensorFill", [], [tensor.name],
            arg=[
                utils.MakeArgument("shape", list(tensor.dims)),
                utils.MakeArgument("values", tensor.float_data)])
        init_net.op.extend([op])
    return init_net


def BaseTranslate(layer, caffe2_type):
    """A simple translate interface that maps the layer input and output."""
    caffe2_op = caffe2_pb2.OperatorDef()
    caffe2_op.type = caffe2_type
    caffe2_op.input.extend(layer.bottom)
    caffe2_op.output.extend(layer.top)
    return caffe2_op


def AddArgument(op, key, value):
    """Makes an argument based on the value type."""
    op.arg.extend([utils.MakeArgument(key, value)])

################################################################################
# Common translators for layers.
################################################################################


@TranslatorRegistry.Register("Input")
def TranslateInput(layer, pretrained_blobs, is_test):
    return [], []


@TranslatorRegistry.Register("Data")
def TranslateData(layer, pretrained_blobs, is_test):
    return [], []


# A function used in convolution, pooling and deconvolution to deal with
# conv pool specific parameters.
def _TranslateStridePadKernelHelper(param, caffe_op):
    try:
        if (len(param.stride) > 1 or len(param.kernel_size) > 1 or
                len(param.pad) > 1):
            raise NotImplementedError(
                "Translator currently does not support non-conventional "
                "pad/kernel/stride settings."
            )
        stride = param.stride[0] if len(param.stride) else 1
        pad = param.pad[0] if len(param.pad) else 0
        kernel = param.kernel_size[0] if len(param.kernel_size) else 0
    except TypeError:
        # This catches the case of a PoolingParameter, in which case we are
        # having non-repeating pad, stride and kernel.
        stride = param.stride
        pad = param.pad
        kernel = param.kernel_size
    # Get stride
    if param.HasField("stride_h") or param.HasField("stride_w"):
        AddArgument(caffe_op, "stride_h", param.stride_h)
        AddArgument(caffe_op, "stride_w", param.stride_w)
    else:
        AddArgument(caffe_op, "stride", stride)
    # Get pad
    if param.HasField("pad_h") or param.HasField("pad_w"):
        if param.pad_h == param.pad_w:
            AddArgument(caffe_op, "pad", param.pad_h)
        else:
            AddArgument(caffe_op, "pad_t", param.pad_h)
            AddArgument(caffe_op, "pad_b", param.pad_h)
            AddArgument(caffe_op, "pad_l", param.pad_w)
            AddArgument(caffe_op, "pad_r", param.pad_w)
    else:
        AddArgument(caffe_op, "pad", pad)
    # Get kernel
    if param.HasField("kernel_h") or param.HasField("kernel_w"):
        AddArgument(caffe_op, "kernel_h", param.kernel_h)
        AddArgument(caffe_op, "kernel_w", param.kernel_w)
    else:
        AddArgument(caffe_op, "kernel", kernel)


@TranslatorRegistry.Register("Convolution")
def TranslateConv(layer, pretrained_blobs, is_test):
    param = layer.convolution_param
    if param.group > 1:
        return TranslateConvWithGroups(layer, pretrained_blobs, is_test)
    # If there is no odd things, we will basically translate it to a standard
    # caffe2 op.
    caffe_op = BaseTranslate(layer, "Conv")
    output = caffe_op.output[0]
    caffe_op.input.extend([output + '_w', output + '_b'])
    _TranslateStridePadKernelHelper(param, caffe_op)
    weight = utils.NumpyArrayToCaffe2Tensor(pretrained_blobs[0], output + '_w')
    bias = utils.NumpyArrayToCaffe2Tensor(
        pretrained_blobs[1].flatten(), output + '_b'
    )
    return caffe_op, [weight, bias]


def TranslateConvWithGroups(layer, pretrained_blobs, is_test):
    print(
        "Legacy warning: convolution with groups seem to be less and less " +
        "popular, so we no longer have it as a first-class citizen op. " +
        "Instead, we will simulate it with depth split followed by conv " +
        "followed by depth concat."
    )
    caffe_ops = []
    caffe_params = []
    param = layer.convolution_param
    weight, bias = pretrained_blobs
    bias = bias.flatten()
    n, c, h, w = weight.shape
    g = param.group  # group
    od = int(n / g)  # output dimension
    if (od * g != n):
        # This should not happen: n should always be divisible by g.
        raise ValueError("This should not happen.")
    output = layer.top[0]
    # first, depth_split
    depth_split_op = core.CreateOperator(
        "DepthSplit",
        str(layer.bottom[0]),
        ['_' + output + '_gconv_split_' + str(i) for i in range(g)],
        split=[c for i in range(g)],
        order="NCHW"
    )
    caffe_ops.append(depth_split_op)
    # second, convolutions
    for i in range(g):
        # convolution layer i
        this_weight = utils.NumpyArrayToCaffe2Tensor(
            weight[i * od:(i + 1) * od], output + '_gconv_' + str(i) + '_w'
        )
        this_bias = utils.NumpyArrayToCaffe2Tensor(
            bias[i * od:(i + 1) * od], output + '_gconv_' + str(i) + '_b'
        )
        conv_op = core.CreateOperator(
            "Conv",
            [depth_split_op.output[i], this_weight.name, this_bias.name],
            ['_' + output + '_gconv_conv_' + str(i)],
            order="NCHW"
        )
        _TranslateStridePadKernelHelper(param, conv_op)
        caffe_ops.append(conv_op)
        caffe_params.extend([this_weight, this_bias])
    # third, depth concat
    depth_concat_op = core.CreateOperator(
        "Concat",
        ['_' + output + '_gconv_conv_' + str(i) for i in range(g)],
        [output, '_' + output + '_gconv_concat_dims'],
        order="NCHW"
    )
    caffe_ops.append(depth_concat_op)
    return caffe_ops, caffe_params


@TranslatorRegistry.Register("Deconvolution")
def TranslateDeconv(layer, pretrained_blobs, is_test):
    param = layer.convolution_param
    if param.group > 1:
        raise NotImplementedError(
            "Translator currently does not support group deconvolution."
        )
    caffe_op = BaseTranslate(layer, "ConvTranspose")
    output = caffe_op.output[0]
    _TranslateStridePadKernelHelper(param, caffe_op)
    caffe_op.input.extend([output + '_w', output + '_b'])
    AddArgument(caffe_op, "order", "NCHW")
    weight = utils.NumpyArrayToCaffe2Tensor(pretrained_blobs[0], output + '_w')
    bias = utils.NumpyArrayToCaffe2Tensor(
        pretrained_blobs[1].flatten(), output + '_b'
    )
    return caffe_op, [weight, bias]


@TranslatorRegistry.Register("ReLU")
def TranslateRelu(layer, pretrained_blobs, is_test):
    return BaseTranslate(layer, "Relu"), []


@TranslatorRegistry.Register("Pooling")
def TranslatePool(layer, pretrained_blobs, is_test):
    param = layer.pooling_param
    if param.pool == caffe_pb2.PoolingParameter.MAX:
        caffe_op = BaseTranslate(layer, "MaxPool")
    elif param.pool == caffe_pb2.PoolingParameter.AVE:
        caffe_op = BaseTranslate(layer, "AveragePool")
    _TranslateStridePadKernelHelper(param, caffe_op)
    AddArgument(caffe_op, "order", "NCHW")
    AddArgument(caffe_op, "legacy_pad",
                caffe2_legacy_pb2.CAFFE_LEGACY_POOLING)
    if param.global_pooling:
        AddArgument(caffe_op, "global_pooling", 1)
    return caffe_op, []


@TranslatorRegistry.Register("LRN")
def TranslateLRN(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "LRN")
    caffe_op.output.extend(['_' + caffe_op.output[0] + '_scale'])
    param = layer.lrn_param
    if param.norm_region != caffe_pb2.LRNParameter.ACROSS_CHANNELS:
        raise ValueError(
            "Does not support norm region other than across channels.")
    AddArgument(caffe_op, "size", int(param.local_size))
    AddArgument(caffe_op, "alpha", float(param.alpha))
    AddArgument(caffe_op, "beta", float(param.beta))
    AddArgument(caffe_op, "bias", float(param.k))
    AddArgument(caffe_op, "order", "NCHW")
    return caffe_op, []


@TranslatorRegistry.Register("InnerProduct")
def TranslateInnerProduct(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "FC")
    output = caffe_op.output[0]
    caffe_op.input.extend([output + '_w', output + '_b'])
    weight = utils.NumpyArrayToCaffe2Tensor(
        pretrained_blobs[0][0, 0], output + '_w'
    )
    bias = utils.NumpyArrayToCaffe2Tensor(
        pretrained_blobs[1].flatten(), output + '_b'
    )
    return caffe_op, [weight, bias]


@TranslatorRegistry.Register("Dropout")
def TranslateDropout(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "Dropout")
    caffe_op.output.extend(['_' + caffe_op.output[0] + '_mask'])
    param = layer.dropout_param
    AddArgument(caffe_op, "ratio", param.dropout_ratio)
    if (is_test):
        AddArgument(caffe_op, "is_test", 1)
    return caffe_op, []


@TranslatorRegistry.Register("Softmax")
def TranslateSoftmax(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "Softmax")
    return caffe_op, []


@TranslatorRegistry.Register("SoftmaxWithLoss")
def TranslateSoftmaxWithLoss(layer, pretrained_blobs, is_test):
    softmax_op = core.CreateOperator(
        "Softmax", [layer.bottom[0]],
        layer.bottom[0] + "_translator_autogen_softmax")
    xent_op = core.CreateOperator(
        "LabelCrossEntropy",
        [softmax_op.output[0], layer.bottom[1]],
        layer.bottom[0] + "_translator_autogen_xent")
    loss_op = core.CreateOperator(
        "AveragedLoss",
        xent_op.output[0],
        layer.top[0])
    return [softmax_op, xent_op, loss_op], []


@TranslatorRegistry.Register("Accuracy")
def TranslateAccuracy(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "Accuracy")
    if layer.accuracy_param.top_k != 1:
        print("Warning: Translation does not support Accuracy layers top_k >1.")
    return caffe_op, []


@TranslatorRegistry.Register("Concat")
def TranslateConcat(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "Concat")
    caffe_op.output.extend(['_' + caffe_op.output[0] + '_dims'])
    AddArgument(caffe_op, "order", "NCHW")
    return caffe_op, []


@TranslatorRegistry.Register("TanH")
def TranslateTanH(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "Tanh")
    return caffe_op, []


@TranslatorRegistry.Register("InstanceNorm")
def TranslateInstanceNorm(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "InstanceNorm")
    output = caffe_op.output[0]
    weight = utils.NumpyArrayToCaffe2Tensor(
        pretrained_blobs[0].flatten(), output + '_w')
    bias = utils.NumpyArrayToCaffe2Tensor(
        pretrained_blobs[1].flatten(), output + '_b')
    caffe_op.input.extend([output + '_w', output + '_b'])
    AddArgument(caffe_op, "order", "NCHW")
    return caffe_op, [weight, bias]


@TranslatorRegistry.Register("Eltwise")
def TranslateElementWise(layer, pretrained_blobs, is_test):
    param = layer.eltwise_param
    # TODO(jiayq): if we have a protobuf that uses this, lift this constraint
    # and verify that we can correctly translate.
    if len(param.coeff) or param.operation != 1:
        raise RuntimeError("This eltwise layer is not yet supported.")
    caffe_op = BaseTranslate(layer, "Sum")
    return caffe_op, []


@TranslatorRegistry.Register("Scale")
def TranslateScale(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "Mul")
    scale_param = layer.scale_param
    AddArgument(caffe_op, "axis", scale_param.axis)
    AddArgument(caffe_op, "broadcast", True)
    if len(caffe_op.input) == 1:
        # the scale parameter is in pretrained blobs
        if scale_param.num_axes != 1:
            raise RuntimeError("This path has not been verified yet.")
        output = caffe_op.output[0]
        caffe_op.input.append(output + '_w')
        weight = utils.NumpyArrayToCaffe2Tensor(
            pretrained_blobs[0].flatten(), output + '_w')
        return caffe_op, [weight]
    elif len(caffe_op.input) == 2:
        # TODO(jiayq): find a protobuf that uses this and verify.
        raise RuntimeError("This path has not been verified yet.")
    else:
        raise RuntimeError("Unexpected number of inputs.")


@TranslatorRegistry.Register("Reshape")
def TranslateReshape(layer, pretrained_blobs, is_test):
    caffe_op = BaseTranslate(layer, "Reshape")
    caffe_op.output.append("_" + caffe_op.input[0] + "_dims")
    reshape_param = layer.reshape_param
    AddArgument(caffe_op, 'shape', reshape_param.shape.dim)
    return caffe_op, []