caffe_decoder.py

#   Copyright (c) 2019  PaddlePaddle Authors. All Rights Reserved.
#
# 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.

import os
import sys
from google.protobuf import text_format
import numpy as np
from x2paddle.core.graph import GraphNode, Graph
from x2paddle.core.fluid_code import FluidCode
from x2paddle.op_mapper import caffe_shape


class CaffeResolver(object):
    def __init__(self, caffe_proto):
        self.caffe_proto = caffe_proto
        self.import_caffe()

    def import_caffepb(self):
        if self.caffe_proto is None:
            from x2paddle.decoder import caffe_pb2
            out = caffe_pb2
        else:
            if not os.path.isfile(self.caffe_proto):
                raise Exception(
                    "The .py file compiled by caffe.proto is not exist.")
            (filepath,
             tempfilename) = os.path.split(os.path.abspath(self.caffe_proto))
            (filename, extension) = os.path.splitext(tempfilename)
            sys.path.append(filepath)
            out = __import__(filename)
        return out

    def import_caffe(self):
        self.caffepb = self.import_caffepb()
        self.NetParameter = self.caffepb.NetParameter


class CaffeGraphNode(GraphNode):
    def __init__(self, layer, layer_name=None):
        if layer_name is None:
            super(CaffeGraphNode,
                  self).__init__(layer,
                                 layer.name.replace('/', '_').replace('-', '_'))
        else:
            super(CaffeGraphNode,
                  self).__init__(layer,
                                 layer_name.replace('/', '_').replace('-', '_'))
        self.layer_type = layer.type
        self.fluid_code = FluidCode()
        self.data = None

    def set_params(self, params):
        self.data = params

    def set_output_shape(self, input_shape, is_input=True):
        func_name = 'shape_' + self.layer_type.lower()
        if is_input:
            self.output_shape = getattr(caffe_shape, func_name)(self.layer,
                                                                input_shape)
        else:
            self.output_shape = input_shape

    def set_input_shape(self, input_shape):
        self.input_shape = input_shape


class CaffeGraph(Graph):
    def __init__(self, model, params):
        self.params = params
        super(CaffeGraph, self).__init__(model)

    def filter_layers(self, layers):
        '''Filter out layers based on the current phase.'''
        phase_map = {0: 'train', 1: 'test'}
        filtered_layer_names = set()
        filtered_layers = []
        for layer in layers:
            phase = 'test'
            if len(layer.include):
                phase = phase_map[layer.include[0].phase]
            if len(layer.exclude):
                phase = phase_map[1 - layer.include[0].phase]
            exclude = (phase != 'test')
            # Dropout layers appear in a fair number of Caffe
            # test-time networks. These are just ignored. We'll
            # filter them out here.
            if (not exclude) and (phase == 'test'):
                exclude = (layer.type == 'Dropout')
            if not exclude:
                filtered_layers.append(layer)
                # Guard against dupes.
                assert layer.name not in filtered_layer_names
                filtered_layer_names.add(layer.name)
            else:
                print('The filter layer:' + layer.name)
        return filtered_layers

    def build(self):
        layers = self.model.layers or self.model.layer
        layers = self.filter_layers(layers)

        top_layer = {}
        for layer in layers:
            self.node_map[layer.name] = CaffeGraphNode(layer)
            for in_name in layer.bottom:
                if in_name in top_layer:
                    self.connect(top_layer[in_name][-1], layer.name)
                else:
                    raise Exception(
                        'input[{}] of node[{}] does not exist in node_map'.
                        format(in_name, layer.name))
            for out_name in layer.top:
                if out_name not in top_layer:
                    top_layer[out_name] = [layer.name]
                else:
                    top_layer[out_name].append(layer.name)

        for layer_name, data in self.params:
            if layer_name in self.node_map:
                node = self.node_map[layer_name]
                node.set_params(data)
            else:
                print('Ignoring parameters for non-existent layer: %s' % \
                       layer_name)

        super(CaffeGraph, self).build()

    def get_bottom_node(self, node, idx=0, copy=False):
        input_node_name = node.inputs[idx]
        assert input_node_name in self.node_map, 'The {} isn\'t a valid node'.format(
            name)
        input_node = self.node_map[input_node_name]
        if len(input_node.layer.top) > 1:
            need_idx = list(input_node.layer.top).index(node.layer.bottom[idx])
            name = input_node_name + ':' + str(need_idx)
        else:
            name = input_node_name
        return self.get_node(name, copy=copy)


class CaffeDecoder(object):
    def __init__(self, proto_path, model_path, caffe_proto):
        self.proto_path = proto_path
        self.model_path = model_path

        self.resolver = CaffeResolver(caffe_proto=caffe_proto)
        self.net = self.resolver.NetParameter()
        with open(proto_path, 'rb') as proto_file:
            proto_str = self.old2new(proto_file)
            text_format.Merge(proto_str, self.net)

        self.load_using_pb()
        self.caffe_graph = CaffeGraph(self.net, self.params)
        self.caffe_graph.build()

    def load_using_pb(self):
        data = self.resolver.NetParameter()
        data.MergeFromString(open(self.model_path, 'rb').read())
        pair = lambda layer: (layer.name, self.normalize_pb_data(layer))
        layers = data.layers or data.layer
        self.params = [pair(layer) for layer in layers if layer.blobs]

    def normalize_pb_data(self, layer):
        transformed = []
        for blob in layer.blobs:
            if len(blob.shape.dim):
                dims = blob.shape.dim
                c_o, c_i, h, w = map(int, [1] * (4 - len(dims)) + list(dims))
            else:
                c_o = blob.num
                c_i = blob.channels
                h = blob.height
                w = blob.width
            data = np.array(blob.data, dtype=np.float32).reshape(c_o, c_i, h, w)
            transformed.append(data)
        return transformed

    def old2new(self, proto_file):
        part1_str = ''
        part2_str = ''
        part3_str = ''
        is_input = False
        dims = []
        line = proto_file.readline()
        print('Check if it is a new style of caffe...')
        while line:
            l_str = bytes.decode(line)
            if l_str.replace(' ', '').startswith('input:'):
                part2_str += 'layer {\n'
                part2_str += (
                    '  name: ' +
                    l_str.strip().replace(' ', '').split('input:')[-1] + '\n')
                part2_str += '  type: \"Input\"\n'
                part2_str += (
                    '  top: ' +
                    l_str.strip().replace(' ', '').split('input:')[-1] + '\n')
                is_input = True
                line = proto_file.readline()
                continue
            elif l_str.replace(' ', '').startswith('input_dim:'):
                dims.append(
                    int(l_str.strip().replace(' ', '').split('input_dim:')[-1]))
                if len(dims) == 4:
                    part2_str += '  input_param { shape: { dim: ' + str(dims[0]) + \
                                               ' dim: ' + str(dims[1]) + \
                                               ' dim: ' + str(dims[2]) + \
                                               ' dim: ' + str(dims[3]) + ' } }\n'
                    dims = []
                    part2_str += '}\n'
                line = proto_file.readline()
                if bytes.decode(line).replace(' ', '').startswith('}'):
                    line = proto_file.readline()
                continue
            elif l_str.replace(' ', '').startswith('input_shape'):
                part2_str += l_str.replace('input_shape',
                                           'input_param { shape: ')
                l_str = bytes.decode(proto_file.readline())
                while l_str:
                    if '}' in l_str:
                        part2_str += l_str + '\n}\n}'
                        break
                    else:
                        part2_str += l_str
                    l_str = bytes.decode(proto_file.readline())
                line = proto_file.readline()
                continue
            if not is_input:
                part1_str += bytes.decode(line)
            else:
                part3_str += bytes.decode(line)
            line = proto_file.readline()
        out = part1_str + part2_str + part3_str
        layer_str = 'layer{'
        part = out.split(layer_str)
        if len(part) == 1:
            layer_str = 'layer {'
            part = out.split(layer_str)
        for i in range(len(part)):
            if part[i].strip().replace(' ', '') == '' or part[i].count(':') > 1:
                continue
            out = out.replace(layer_str + part[i], part[i].replace(' ', ''))
        return str.encode(out)