model_convert_utils.py

# Copyright (c) 2017-present, Facebook, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################

'''Helper functions for model conversion to pb'''

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

from functools import wraps
import copy
import numpy as np

from caffe2.python import core, workspace
from caffe2.proto import caffe2_pb2


class OpFilter(object):
    def __init__(self, **kwargs):
        self.type = None
        self.type_in = None
        self.inputs = None
        self.outputs = None
        self.input_has = None
        self.output_has = None
        self.cond = None
        self.reverse = False

        assert all([x in self.__dict__ for x in kwargs])
        self.__dict__.update(kwargs)

    def check(self, op):
        ret = self.reverse
        if self.type and op.type != self.type:
            return ret
        if self.type_in and op.type not in self.type_in:
            return ret
        if self.inputs and set(op.input) != set(self.inputs):
            return ret
        if self.outputs and set(op.output) != set(self.outputs):
            return ret
        if self.input_has and self.input_has not in op.input:
            return ret
        if self.output_has and self.output_has not in op.output:
            return ret
        if self.cond is not None and not self.cond:
            return ret
        return not ret


def filter_op(op, **kwargs):
    ''' Returns true if passed all checks '''
    return OpFilter(**kwargs).check(op)


def op_filter(**filter_args):
    ''' Returns None if no condition is satisfied '''
    def actual_decorator(f):
        @wraps(f)
        def wrapper(op, **params):
            if not filter_op(op, **filter_args):
                return None
            return f(op, **params)
        return wrapper
    return actual_decorator


def op_func_chain(convert_func_list):
    ''' Run funcs one by one until func return is not None '''
    assert isinstance(convert_func_list, list)

    def _chain(op):
        for x in convert_func_list:
            ret = x(op)
            if ret is not None:
                return ret
        return None

    return _chain


def convert_op_in_ops(ops_ref, func_or_list):
    func = func_or_list
    if isinstance(func_or_list, list):
        func = op_func_chain(func_or_list)
    ops = [op for op in ops_ref]
    converted_ops = []
    for op in ops:
        new_ops = func(op)
        if new_ops is not None and not isinstance(new_ops, list):
            new_ops = [new_ops]
        converted_ops.extend(new_ops if new_ops is not None else [op])
    del ops_ref[:]
    # ops_ref maybe of type RepeatedCompositeFieldContainer
    # which does not have append()
    ops_ref.extend(converted_ops)


def convert_op_in_proto(proto, func_or_list):
    convert_op_in_ops(proto.op, func_or_list)


def get_op_arg(op, arg_name):
    for x in op.arg:
        if x.name == arg_name:
            return x
    return None


def get_op_arg_valf(op, arg_name, default_val):
    arg = get_op_arg(op, arg_name)
    return arg.f if arg is not None else default_val


def update_mobile_engines(net):
    for op in net.op:
        if op.type == "Conv":
            op.engine = "NNPACK"
        if op.type == "ConvTranspose":
            op.engine = "BLOCK"


def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    from itertools import tee
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


def blob_uses(net, blob):
    u = []
    for i, op in enumerate(net.op):
        if blob in op.input or blob in op.control_input:
            u.append(i)
    return u


def fuse_first_affine(net, params, removed_tensors):
    net = copy.deepcopy(net)
    params = copy.deepcopy(params)

    for ((i, current), (j, next_)) in pairwise(enumerate(net.op)):
        if next_.input[0] != current.output[0]:
            continue

        if current.type not in ("Conv", "ConvTranspose") \
           or next_.type != "AffineChannel":
            continue
        if current.output[0] != next_.output[0] and \
                len(blob_uses(net, current.output[0])) != 1:
            # Can't fuse if more than one user unless AffineChannel is inplace
            continue

        # else, can fuse
        conv = current
        affine = next_
        fused_conv = copy.deepcopy(conv)
        fused_conv.output[0] = affine.output[0]
        conv_weight = params[conv.input[1]]
        conv_has_bias = len(conv.input) > 2
        conv_bias = params[conv.input[2]] if conv_has_bias else 0

        A = params[affine.input[1]]
        B = params[affine.input[2]]

        # Thus, can just have the affine transform
        # X * A + B
        # where
        # A = bn_scale * 1.0 / (sqrt(running_var + eps))
        # B =  (bias - running_mean * (1.0 / sqrt(running_var + eps))
        # * bn_scale)

        # This identify should hold if we have correctly fused
        # np.testing.assert_array_equal(
        #     params[conv.output[0]] * A + B,
        #     params[bn.output[0]])

        # Now, we have that the computation made is the following:
        # ((X `conv` W) + b) * A + B
        # Then, we can simply fuse this as follows:
        # (X `conv` (W * A)) + b * A + B
        # which is simply
        # (X `conv` Q) + C
        # where

        # Q = W * A
        # C = b * A + B

        # For ConvTranspose, from the view of convolutions as a
        # Toepeliz multiplication, we have W_ = W^T, so the weights
        # are laid out as (R, S, K, K) (vs (S, R, K, K) for a Conv),
        # so the weights broadcast slightly differently. Remember, our
        # BN scale 'B' is of size (S,)

        A_ = A.reshape(-1, 1, 1, 1) if conv.type == "Conv" else \
            A.reshape(1, -1, 1, 1)

        C = conv_bias * A + B
        Q = conv_weight * A_

        assert params[conv.input[1]].shape == Q.shape

        params[conv.input[1]] = Q
        if conv_has_bias:
            assert params[conv.input[2]].shape == C.shape
            params[conv.input[2]] = C
        else:
            # make af_bias to be bias of the conv layer
            fused_conv.input.append(affine.input[2])
            params[affine.input[2]] = B

        new_ops = net.op[:i] + [fused_conv] + net.op[j + 1:]
        del net.op[:]
        if conv_has_bias:
            del params[affine.input[2]]
            removed_tensors.append(affine.input[2])
        removed_tensors.append(affine.input[1])
        del params[affine.input[1]]
        net.op.extend(new_ops)
        break
    return net, params, removed_tensors


def fuse_affine(net, params, ignore_failure):
    # Run until we hit a fixed point
    removed_tensors = []
    while True:
        (next_net, next_params, removed_tensors) = \
            fuse_first_affine(net, params, removed_tensors)
        if len(next_net.op) == len(net.op):
            if (
                any(op.type == "AffineChannel" for op in next_net.op) and
                not ignore_failure
            ):
                raise Exception(
                    "Model contains AffineChannel op after fusion: %s", next_net)
            return (next_net, next_params, removed_tensors)
        net, params, removed_tensors = (next_net, next_params, removed_tensors)


def fuse_net(fuse_func, net, blobs, ignore_failure=False):
    is_core_net = isinstance(net, core.Net)
    if is_core_net:
        net = net.Proto()

    net, params, removed_tensors = fuse_func(net, blobs, ignore_failure)
    for rt in removed_tensors:
        net.external_input.remove(rt)

    if is_core_net:
        net = core.Net(net)

    return net, params


def fuse_net_affine(net, blobs):
    return fuse_net(fuse_affine, net, blobs)


def add_tensor(net, name, blob):
    ''' Create an operator to store the tensor 'blob',
        run the operator to put the blob to workspace.
        uint8 is stored as an array of string with one element.
    '''
    kTypeNameMapper = {
        np.dtype('float32'): "GivenTensorFill",
        np.dtype('int32'): "GivenTensorIntFill",
        np.dtype('int64'): "GivenTensorInt64Fill",
        np.dtype('uint8'): "GivenTensorStringFill",
    }

    shape = blob.shape
    values = blob
    # pass array of uint8 as a string to save storage
    # storing uint8_t has a large overhead for now
    if blob.dtype == np.dtype('uint8'):
        shape = [1]
        values = [str(blob.data)]

    op = core.CreateOperator(
        kTypeNameMapper[blob.dtype],
        [], [name],
        shape=shape,
        values=values,
        # arg=[
        #     putils.MakeArgument("shape", shape),
        #     putils.MakeArgument("values", values),
        # ]
    )
    net.op.extend([op])


def gen_init_net_from_blobs(blobs, blobs_to_use=None, excluded_blobs=None):
    ''' Generate an initialization net based on a blob dict '''
    ret = caffe2_pb2.NetDef()
    if blobs_to_use is None:
        blobs_to_use = {x for x in blobs}
    else:
        blobs_to_use = copy.deepcopy(blobs_to_use)
    if excluded_blobs is not None:
        blobs_to_use = [x for x in blobs_to_use if x not in excluded_blobs]
    for name in blobs_to_use:
        blob = blobs[name]
        if isinstance(blob, str):
            print('Blob {} with type {} is not supported in generating init net,'
                  ' skipped.'.format(name, type(blob)))
            continue
        add_tensor(ret, name, blob)

    return ret


def get_ws_blobs(blob_names=None):
    ''' Get blobs in 'blob_names' in the default workspace,
        get all blobs if blob_names is None '''
    blobs = {}
    if blob_names is None:
        blob_names = workspace.Blobs()
    blobs = {x: workspace.FetchBlob(x) for x in blob_names}

    return blobs


def get_device_option_cpu():
    device_option = core.DeviceOption(caffe2_pb2.CPU)
    return device_option


def get_device_option_cuda(gpu_id=0):
    device_option = caffe2_pb2.DeviceOption()
    device_option.device_type = caffe2_pb2.CUDA
    device_option.cuda_gpu_id = gpu_id
    return device_option


def create_input_blobs_for_net(net_def):
    for op in net_def.op:
        for blob_in in op.input:
            if not workspace.HasBlob(blob_in):
                workspace.CreateBlob(blob_in)


def compare_model(model1_func, model2_func, test_image, check_blobs):
    ''' model_func(test_image, check_blobs)
    '''
    cb1, cb2 = check_blobs, check_blobs
    if isinstance(check_blobs, dict):
        cb1 = check_blobs.keys()
        cb2 = check_blobs.values()
    print('Running the first model...')
    res1 = model1_func(test_image, check_blobs)
    print('Running the second model...')
    res2 = model2_func(test_image, check_blobs)
    for idx in range(len(cb1)):
        print('Checking {} -> {}...'.format(cb1[idx], cb2[idx]))
        n1, n2 = cb1[idx], cb2[idx]
        r1 = res1[n1] if n1 in res1 else None
        r2 = res2[n2] if n2 in res2 else None
        assert r1 is not None, \
            "Blob {} in model1 is None".format(n1)
        assert r2 is not None, \
            "Blob {} in model2 is None".format(n2)
        assert r1.shape == r2.shape, \
            "Blob {} and {} shape mismatched: {} vs {}".format(
                n1, n2, r1.shape, r2.shape)

        np.testing.assert_array_almost_equal(
            r1, r2, decimal=3,
            err_msg='{} and {} not matched. Max diff: {}'.format(
                n1, n2, np.amax(np.absolute(r1 - r2))))

    return True


# graph_name could not contain word 'graph'
def save_graph(net, file_name, graph_name="net", op_only=True):
    from caffe2.python import net_drawer
    graph = None
    ops = net.op
    if not op_only:
        graph = net_drawer.GetPydotGraph(
            ops, graph_name,
            rankdir="TB")
    else:
        graph = net_drawer.GetPydotGraphMinimal(
            ops, graph_name,
            rankdir="TB", minimal_dependency=True)

    try:
        graph.write_png(file_name)
    except Exception as e:
        print('Error when writing graph to image {}'.format(e))