diff --git a/mmdnn/conversion/caffe/graph.py b/mmdnn/conversion/caffe/graph.py index b91a89cf..28d9e103 100644 --- a/mmdnn/conversion/caffe/graph.py +++ b/mmdnn/conversion/caffe/graph.py @@ -1,451 +1,460 @@ -from collections import namedtuple -from functools import reduce -from google.protobuf import text_format -from copy import deepcopy -import numbers -import os -import tempfile - -from mmdnn.conversion.caffe.mapper import get_handler_name -from mmdnn.conversion.caffe.resolver import get_caffe_resolver, has_pycaffe -from mmdnn.conversion.caffe.shape import * -from mmdnn.conversion.caffe.errors import print_stderr, ConversionError - - -layer_num_to_name = { - 0: 'None', - 1: 'Accuracy', - 2: 'BNLL', - 3: 'Concat', - 4: 'Convolution', - 5: 'Data', - 6: 'Dropout', - 7: 'EuclideanLoss', - 8: 'Flatten', - 9: 'HDF5Data', - 10: 'HDF5Output', - 11: 'Im2col', - 12: 'ImageData', - 13: 'InfogainLoss', - 14: 'InnerProduct', - 15: 'LRN', - 16: 'MultinomialLogisticLoss', - 17: 'Pooling', - 18: 'ReLU', - 19: 'Sigmoid', - 20: 'Softmax', - 21: 'SoftmaxWithLoss', - 22: 'Split', - 23: 'TanH', - 24: 'WindowData', - 25: 'Eltwise', - 26: 'Power', - 27: 'SigmoidCrossEntropyLoss', - 28: 'HingeLoss', - 29: 'MemoryData', - 30: 'ArgMax', - 31: 'Threshold', - 32: 'DummyData', - 33: 'Slice', - 34: 'MVN', - 35: 'AbsVal', - 36: 'Silence', - 37: 'ContrastiveLoss', - 38: 'Exp', - 39: 'Deconvolution', - 40: 'PReLU' - } - -LAYER_DESCRIPTORS = { - # Caffe Types - 'AbsVal': shape_identity, - 'Accuracy': shape_scalar, - 'ArgMax': shape_not_implemented, - 'BatchNorm': shape_identity, - 'BNLL': shape_not_implemented, - 'Concat': shape_concat, - 'ContrastiveLoss': shape_scalar, - 'Convolution': shape_convolution, - 'Crop': shape_not_implemented, - 'Deconvolution': shape_deconvolution, - 'Data': shape_data, - 'Dropout': shape_identity, - 'DummyData': shape_data, - 'EuclideanLoss': shape_scalar, - 'Eltwise': shape_identity, - 'Exp': shape_identity, - 'Flatten': shape_flatten, - 'HDF5Data': shape_data, - 'HDF5Output': shape_identity, - 'HingeLoss': shape_scalar, - 'Im2col': shape_not_implemented, - 'ImageData': shape_data, - 'InfogainLoss': shape_scalar, - 'InnerProduct': shape_inner_product, - 'Input': shape_data, - 'LRN': shape_identity, - 'MemoryData': shape_mem_data, - 'MultinomialLogisticLoss': shape_scalar, - 'MVN': shape_not_implemented, - 'Pooling': shape_pool, - 'Power': shape_identity, - 'ReLU': shape_identity, - 'Scale': shape_identity, - 'Sigmoid': shape_identity, - 'SigmoidCrossEntropyLoss': shape_scalar, - 'Silence': shape_identity, - 'Softmax': shape_identity, - 'SoftmaxWithLoss': shape_scalar, - 'Split': shape_split, - 'Slice': shape_not_implemented, - 'TanH': shape_identity, - 'WindowData': shape_not_implemented, - 'Threshold': shape_identity, - 'Reshape' : shape_reshape, - 'ResizeBilinear': shape_reshape, - 'PReLU' : shape_identity - } - -LAYER_TYPES = LAYER_DESCRIPTORS.keys() - -LayerType = type('LayerType', (), {t : t for t in LAYER_TYPES}) - -KernelParameters = namedtuple('KernelParameters', ['global_pooling', 'k_h', 'k_w', 's_h', 's_w', 'p_h', 'p_w']) - -class NodeKind(LayerType): - - @staticmethod - def map_raw_kind(node_kind): - if isinstance(node_kind, int): - node_kind = layer_num_to_name[node_kind] - else: - node_kind = str(node_kind) - if node_kind in LAYER_TYPES: - return node_kind - return None - - @staticmethod - def compute_output_shape(node): - try: - return LAYER_DESCRIPTORS[node.kind](node) - except NotImplementedError: - raise ConversionError('Output shape computation not implemented for type: %s' % node.kind) - -LAYER_IN_TRAIN_PROTO = [NodeKind.ImageData, NodeKind.Data, NodeKind.HDF5Data, NodeKind.HDF5Output, NodeKind.WindowData, NodeKind.DummyData, NodeKind.MemoryData] - -class CaffeNode(object): - def __init__(self, name, kind, layer=None): - self.name = name - self.kind = kind - self.layer = layer - self.parents = [] - self.children = [] - self.data = None - self.output = [] - self.output_shape = None - self.metadata = {} - - def add_parent(self, parent_node, from_output): - assert parent_node not in self.parents - self.parents.append((parent_node, from_output)) - if self not in parent_node.children: - parent_node.children.append(self) - - def get_only_parent(self): - if len(self.parents) != 1: - raise ConversionError('Node (%s) expected to have 1 parent. Found %s.' % (self, len(self.parents))) - return self.parents[0] - - @property - def parameters(self): - if self.layer is not None: - params = get_handler_name(self.kind) - if params == 'deconvolution': - params = 'convolution' - params = '_'.join((params, 'param')) - try: - return getattr(self.layer, params) - except AttributeError: - raise ConversionError('Caffe parameters not found for layer kind: %s' % (self.kind)) - return None - - @staticmethod - def get_kernel_value(scalar, repeated, idx, default=None): - if scalar: - return scalar - if repeated: - if isinstance(repeated, numbers.Number): - return repeated - if len(repeated) == 1: - # Same value applies to all spatial dimensions - return int(repeated[0]) - assert idx < len(repeated) - # Extract the value for the given spatial dimension - return repeated[idx] - if default is None: - raise ValueError('Unable to determine kernel parameter!') - return default - - @property - def kernel_parameters(self): - assert self.kind in (NodeKind.Convolution, NodeKind.Pooling, NodeKind.Deconvolution) - params = self.parameters - global_pooling = hasattr(params, 'global_pooling') and params.global_pooling - if not global_pooling: - k_h = self.get_kernel_value(params.kernel_h, params.kernel_size, 0) - k_w = self.get_kernel_value(params.kernel_w, params.kernel_size, 1) - s_h = self.get_kernel_value(params.stride_h, params.stride, 0, default=1) - s_w = self.get_kernel_value(params.stride_w, params.stride, 1, default=1) - else: - k_h = k_w = 0 - s_h = s_w = 1 - p_h = self.get_kernel_value(params.pad_h, params.pad, 0, default=0) - p_w = self.get_kernel_value(params.pad_w, params.pad, 1, default=0) - return KernelParameters(global_pooling, k_h, k_w, s_h, s_w, p_h, p_w) - - def __str__(self): - return '[%s] %s' % (self.kind, self.name) - - def __repr__(self): - return '%s (0x%x)' %(self.name, id(self)) - - -class CaffeGraph(object): - - def __init__(self, nodes=None, name=None): - self.nodes = nodes or [] - self.node_lut = {node.name: node for node in self.nodes} - self.name = name - self.prototxt = None - - def add_node(self, node): - self.nodes.append(node) - self.node_lut[node.name] = node - - def get_node(self, name): - try: - return self.node_lut[name] - except KeyError: - raise ConversionError('Layer not found: %s' % name) - - def get_input_nodes(self): - return [node for node in self.nodes if len(node.parents) == 0] - - def get_output_nodes(self): - return [node for node in self.nodes if len(node.children) == 0] - - def topologically_sorted(self): - visited = set() - sorted_nodes = [] - def topo_sort_dfs(node, visited, sorted_nodes): - if node in visited: - return - visited.add(node) - for n, idx in node.parents: - topo_sort_dfs(n, visited, sorted_nodes) - sorted_nodes.append(node) - for node in self.nodes: - topo_sort_dfs(node, visited, sorted_nodes) - return sorted_nodes - - def compute_output_shapes(self, model): - sorted_nodes = self.topologically_sorted() - (tmp_handle, tmp_prototxt) = tempfile.mkstemp(suffix=".prototxt") - with open(tmp_prototxt, 'w') as f: - f.write(text_format.MessageToString(model)) - self.prototxt = tmp_prototxt - if has_pycaffe(): - caffe = get_caffe_resolver().caffe - net = caffe.Net(tmp_prototxt, caffe.TEST) - for key, value in net.blobs.items(): - try: - node = self.get_node(key) - dims = list(value.shape) - dims = dims + [1] * (4 - len(dims)) - node.output_shape = TensorShape(*dims) - except: - continue - for node in sorted_nodes: - if node.output_shape is None: - node.output_shape = TensorShape(*NodeKind.compute_output_shape(node)) - os.close(tmp_handle) - else: - for node in sorted_nodes: - node.output_shape = TensorShape(*NodeKind.compute_output_shape(node)) - - # consider rewrite this function to Network.py - def replaced(self, new_nodes): - return CaffeGraph(nodes=new_nodes, name=self.name) - - def transformed(self, transformers): - graph = self - for transformer in transformers: - graph = transformer(graph) - if graph is None: - raise ConversionError('Transformer failed: {}'.format(transformer)) - assert isinstance(graph, CaffeGraph) - return graph - - def __contains__(self, key): - return key in self.node_lut - - def __str__(self): - def get_max_shape(data): - if isinstance(data, dict): - max = 0 - val = None - for k, v in data.items(): - tmp = reduce(lambda x, y: x*y, v.shape) - if tmp > max: - val = v.shape - max = tmp - return val - else: - return data[0].shape - hdr = '{:<20} {:<30} {:>20} {:>20}'.format('Type', 'Name', 'Param', 'Output') - s = [hdr, '-' * 94] - for node in self.topologically_sorted(): - data_shape = get_max_shape(node.data) if node.data else '--' - out_shape = node.output_shape or '--' - s.append('{:<20} {:<30} {!s:>20} {!s:>20}'.format(node.kind, node.name, data_shape, tuple(out_shape))) - return '\n'.join(s) - - -class GraphBuilder(object): - def __init__(self, model_path, input_shape=None, is_train_proto=False, phase='test'): - self.model_path = model_path - self.phase = phase - self.is_train_proto = is_train_proto - self.input_shape = input_shape - self.load() - - def load(self): - self.model = get_caffe_resolver().NetParameter() - with open(self.model_path, 'r') as f: - text_format.Merge(f.read(), self.model) - if self.is_train_proto: - self.process_train_proto() - - def process_train_proto(self): - layers = self.model.layer or self.model.layers - delete_layer = set() - split_op_map = dict() - loss_layers = [layer for layer in layers if NodeKind.map_raw_kind(layer.type) in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss)] - a = [layers.remove(layer) for layer in layers[:] if layer in loss_layers[:-1] or NodeKind.map_raw_kind(layer.type) in LAYER_IN_TRAIN_PROTO] - for layer in layers[:]: - if 'label' in layer.bottom: - if NodeKind.map_raw_kind(layer.type) in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss): - continue - elif NodeKind.map_raw_kind(layer.type) == NodeKind.Split: - for item in layer.top: - delete_layer.add(item) - layers.remove(layer) - elif NodeKind.map_raw_kind(layer.type) == NodeKind.Split: - for item in layer.top: - split_op_map[item] = layer.bottom[0] - layers.remove(layer) - - for layer in layers[:]: - for item in delete_layer: - if item in layer.bottom: - layers.remove(layer) - break - for key, value in split_op_map.items(): - if key in layer.bottom: - layer.bottom.remove(key) - layer.bottom.append(value) - self.model.input.append('data') - self.model.input_dim.extend(self.input_shape) - last_layer = layers[-1] - kind = NodeKind.map_raw_kind(last_layer.type) - if kind in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss): - pred = layers.add() - pred.name = 'prob' - pred.top.append('prob') - pred.bottom.append(last_layer.bottom[0]) - if kind == NodeKind.SoftmaxWithLoss: - pred.type = NodeKind.Softmax if self.model.layer else 20 # competiable with old version caffe proto - elif kind == NodeKind.SigmoidCrossEntropyLoss: - pred.type = NodeKind.Sigmoid if self.model.layer else 19 - layers.remove(last_layer) - - def filter_layers(self, layers): - phase_map = {0: 'train', 1: 'test'} - filtered_layer_names = set() - filtered_layers = [] - for layer in layers: - phase = self.phase - 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 != self.phase) - # 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 == LayerType.Dropout) - if (not exclude): - exclude = (layer.type == LayerType.Silence) - if not exclude: - if layer.name in filtered_layer_names: - for i in range(1, len(filtered_layer_names)): - new_name = layer.name + '_%s' % i - if new_name not in filtered_layer_names: - layer.name = new_name - break - filtered_layer_names.add(layer.name) - filtered_layers.append(layer) - return filtered_layers - - def make_node(self, layer): - kind = NodeKind.map_raw_kind(layer.type) - if kind is None: - # TODO: raise error - pass - node = CaffeNode(layer.name, kind, layer=layer) - node.output.append(layer.name.replace('/', '_')) - node.output.extend(layer.top[1:]) - return node - - def make_input_node(self): - nodes = [CaffeNode(name, NodeKind.Data) for name in self.model.input] - if len(nodes): - input_dim = list(map(int, self.model.input_dim)) - if not input_dim: - if len(self.model.input_shape) > 0: - input_dim = list(map(int, self.model.input_shape[0].dim)) - else: - # TODO: raise error - pass - for node in nodes: - node.output_shape = tuple(input_dim) - node.output.append('data') - return nodes - - def build(self): - layers = self.model.layers or self.model.layer - layers = self.filter_layers(layers) - nodes = self.make_input_node() - nodes += [self.make_node(layer) for layer in layers] - graph = CaffeGraph(nodes=nodes, name=self.model.name) - node_outputs = {} - for idx, layer in enumerate(layers): - node = graph.get_node(layer.name) - for input_name in layer.bottom: - assert input_name != layer.name - parent_node = node_outputs.get(input_name) - if (parent_node is None) or (parent_node==node): - parent_node = graph.get_node(input_name) - if parent_node.layer: - for i, output in enumerate(parent_node.layer.top): - if input_name == output: - node.add_parent(parent_node, i) - else: - node.add_parent(parent_node, 0) - for output_name in layer.top: - if output_name == layer.name: - continue - node_outputs[output_name] = node - graph.compute_output_shapes(self.model) - return graph +<<<<<<< HEAD +from collections import namedtuple +from functools import reduce +from google.protobuf import text_format +from copy import deepcopy +import numbers +import os +import tempfile + +from mmdnn.conversion.caffe.mapper import get_handler_name +from mmdnn.conversion.caffe.resolver import get_caffe_resolver, has_pycaffe +from mmdnn.conversion.caffe.shape import * +from mmdnn.conversion.caffe.errors import print_stderr, ConversionError + + +layer_num_to_name = { + 0: 'None', + 1: 'Accuracy', + 2: 'BNLL', + 3: 'Concat', + 4: 'Convolution', + 5: 'Data', + 6: 'Dropout', + 7: 'EuclideanLoss', + 8: 'Flatten', + 9: 'HDF5Data', + 10: 'HDF5Output', + 11: 'Im2col', + 12: 'ImageData', + 13: 'InfogainLoss', + 14: 'InnerProduct', + 15: 'LRN', + 16: 'MultinomialLogisticLoss', + 17: 'Pooling', + 18: 'ReLU', + 19: 'Sigmoid', + 20: 'Softmax', + 21: 'SoftmaxWithLoss', + 22: 'Split', + 23: 'TanH', + 24: 'WindowData', + 25: 'Eltwise', + 26: 'Power', + 27: 'SigmoidCrossEntropyLoss', + 28: 'HingeLoss', + 29: 'MemoryData', + 30: 'ArgMax', + 31: 'Threshold', + 32: 'DummyData', + 33: 'Slice', + 34: 'MVN', + 35: 'AbsVal', + 36: 'Silence', + 37: 'ContrastiveLoss', + 38: 'Exp', + 39: 'Deconvolution', + 40: 'PReLU' + } + +LAYER_DESCRIPTORS = { + # Caffe Types + 'AbsVal': shape_identity, + 'Accuracy': shape_scalar, + 'ArgMax': shape_not_implemented, + 'BatchNorm': shape_identity, + 'BNLL': shape_not_implemented, + 'Concat': shape_concat, + 'ContrastiveLoss': shape_scalar, + 'Convolution': shape_convolution, + 'Crop': shape_not_implemented, + 'Deconvolution': shape_deconvolution, + 'Data': shape_data, + 'Dropout': shape_identity, + 'DummyData': shape_data, + 'EuclideanLoss': shape_scalar, + 'Eltwise': shape_identity, + 'Exp': shape_identity, + 'Flatten': shape_flatten, + 'HDF5Data': shape_data, + 'HDF5Output': shape_identity, + 'HingeLoss': shape_scalar, + 'Im2col': shape_not_implemented, + 'ImageData': shape_data, + 'InfogainLoss': shape_scalar, + 'InnerProduct': shape_inner_product, + 'Input': shape_data, + 'LRN': shape_identity, + 'MemoryData': shape_mem_data, + 'MultinomialLogisticLoss': shape_scalar, + 'MVN': shape_not_implemented, + 'Pooling': shape_pool, + 'Power': shape_identity, + 'ReLU': shape_identity, + 'Scale': shape_identity, + 'Sigmoid': shape_identity, + 'SigmoidCrossEntropyLoss': shape_scalar, + 'Silence': shape_identity, + 'Softmax': shape_identity, + 'SoftmaxWithLoss': shape_scalar, + 'Split': shape_split, + 'Slice': shape_not_implemented, + 'TanH': shape_identity, + 'WindowData': shape_not_implemented, + 'Threshold': shape_identity, + 'Reshape' : shape_reshape, + 'ResizeBilinear': shape_reshape, + 'PReLU' : shape_identity + } + +LAYER_TYPES = LAYER_DESCRIPTORS.keys() + +LayerType = type('LayerType', (), {t : t for t in LAYER_TYPES}) + +KernelParameters = namedtuple('KernelParameters', ['global_pooling', 'k_h', 'k_w', 's_h', 's_w', 'p_h', 'p_w']) + +class NodeKind(LayerType): + + @staticmethod + def map_raw_kind(node_kind): + if isinstance(node_kind, int): + node_kind = layer_num_to_name[node_kind] + else: + node_kind = str(node_kind) + if node_kind in LAYER_TYPES: + return node_kind + return None + + @staticmethod + def compute_output_shape(node): + try: + return LAYER_DESCRIPTORS[node.kind](node) + except NotImplementedError: + raise ConversionError('Output shape computation not implemented for type: %s' % node.kind) + +LAYER_IN_TRAIN_PROTO = [NodeKind.ImageData, NodeKind.Data, NodeKind.HDF5Data, NodeKind.HDF5Output, NodeKind.WindowData, NodeKind.DummyData, NodeKind.MemoryData] + +class CaffeNode(object): + def __init__(self, name, kind, layer=None): + self.name = name + self.kind = kind + self.layer = layer + self.parents = [] + self.children = [] + self.data = None + self.output = [] + self.output_shape = None + self.metadata = {} + + def add_parent(self, parent_node, from_output): + assert parent_node not in self.parents + self.parents.append((parent_node, from_output)) + if self not in parent_node.children: + parent_node.children.append(self) + + def get_only_parent(self): + if len(self.parents) != 1: + raise ConversionError('Node (%s) expected to have 1 parent. Found %s.' % (self, len(self.parents))) + return self.parents[0] + + @property + def parameters(self): + if self.layer is not None: + params = get_handler_name(self.kind) + if params == 'deconvolution': + params = 'convolution' + params = '_'.join((params, 'param')) + try: + return getattr(self.layer, params) + except AttributeError: + raise ConversionError('Caffe parameters not found for layer kind: %s' % (self.kind)) + return None + + @staticmethod + def get_kernel_value(scalar, repeated, idx, default=None): + if scalar: + return scalar + if repeated: + if isinstance(repeated, numbers.Number): + return repeated + if len(repeated) == 1: + # Same value applies to all spatial dimensions + return int(repeated[0]) + assert idx < len(repeated) + # Extract the value for the given spatial dimension + return repeated[idx] + if default is None: + raise ValueError('Unable to determine kernel parameter!') + return default + + @property + def kernel_parameters(self): + assert self.kind in (NodeKind.Convolution, NodeKind.Pooling, NodeKind.Deconvolution) + params = self.parameters + global_pooling = hasattr(params, 'global_pooling') and params.global_pooling + if not global_pooling: + k_h = self.get_kernel_value(params.kernel_h, params.kernel_size, 0) + k_w = self.get_kernel_value(params.kernel_w, params.kernel_size, 1) + s_h = self.get_kernel_value(params.stride_h, params.stride, 0, default=1) + s_w = self.get_kernel_value(params.stride_w, params.stride, 1, default=1) + else: + k_h = k_w = 0 + s_h = s_w = 1 + p_h = self.get_kernel_value(params.pad_h, params.pad, 0, default=0) + p_w = self.get_kernel_value(params.pad_w, params.pad, 1, default=0) + return KernelParameters(global_pooling, k_h, k_w, s_h, s_w, p_h, p_w) + + def __str__(self): + return '[%s] %s' % (self.kind, self.name) + + def __repr__(self): + return '%s (0x%x)' %(self.name, id(self)) + + +class CaffeGraph(object): + + def __init__(self, nodes=None, name=None): + self.nodes = nodes or [] + self.node_lut = {node.name: node for node in self.nodes} + self.name = name + self.prototxt = None + + def add_node(self, node): + self.nodes.append(node) + self.node_lut[node.name] = node + + def get_node(self, name): + try: + return self.node_lut[name] + except KeyError: + raise ConversionError('Layer not found: %s' % name) + + def get_input_nodes(self): + return [node for node in self.nodes if len(node.parents) == 0] + + def get_output_nodes(self): + return [node for node in self.nodes if len(node.children) == 0] + + def topologically_sorted(self): + visited = set() + sorted_nodes = [] + def topo_sort_dfs(node, visited, sorted_nodes): + if node in visited: + return + visited.add(node) + for n, idx in node.parents: + topo_sort_dfs(n, visited, sorted_nodes) + sorted_nodes.append(node) + for node in self.nodes: + topo_sort_dfs(node, visited, sorted_nodes) + return sorted_nodes + + def compute_output_shapes(self, model): + sorted_nodes = self.topologically_sorted() + (tmp_handle, tmp_prototxt) = tempfile.mkstemp(suffix=".prototxt") + with open(tmp_prototxt, 'w') as f: + f.write(text_format.MessageToString(model)) + self.prototxt = tmp_prototxt + if has_pycaffe(): + caffe = get_caffe_resolver().caffe + net = caffe.Net(tmp_prototxt, caffe.TEST) + for key, value in net.blobs.items(): + try: + node = self.get_node(key) + dims = list(value.shape) + dims = dims + [1] * (4 - len(dims)) + node.output_shape = TensorShape(*dims) + except: + continue + for node in sorted_nodes: + if node.output_shape is None: + top_name = net.top_names.get(node.name) + if top_name is not None: + value = net.blobs.get(top_name[0]) + dims = list(value.shape) + dims = dims + [1] * (4 - len(dims)) + node.output_shape = TensorShape(*dims) + else: + node.output_shape = TensorShape(*NodeKind.compute_output_shape(node)) + + os.close(tmp_handle) + else: + for node in sorted_nodes: + node.output_shape = TensorShape(*NodeKind.compute_output_shape(node)) + + # consider rewrite this function to Network.py + def replaced(self, new_nodes): + return CaffeGraph(nodes=new_nodes, name=self.name) + + def transformed(self, transformers): + graph = self + for transformer in transformers: + graph = transformer(graph) + if graph is None: + raise ConversionError('Transformer failed: {}'.format(transformer)) + assert isinstance(graph, CaffeGraph) + return graph + + def __contains__(self, key): + return key in self.node_lut + + def __str__(self): + def get_max_shape(data): + if isinstance(data, dict): + max = 0 + val = None + for k, v in data.items(): + tmp = reduce(lambda x, y: x*y, v.shape) + if tmp > max: + val = v.shape + max = tmp + return val + else: + return data[0].shape + hdr = '{:<20} {:<30} {:>20} {:>20}'.format('Type', 'Name', 'Param', 'Output') + s = [hdr, '-' * 94] + for node in self.topologically_sorted(): + data_shape = get_max_shape(node.data) if node.data else '--' + out_shape = node.output_shape or '--' + s.append('{:<20} {:<30} {!s:>20} {!s:>20}'.format(node.kind, node.name, data_shape, tuple(out_shape))) + return '\n'.join(s) + + +class GraphBuilder(object): + def __init__(self, model_path, input_shape=None, is_train_proto=False, phase='test'): + self.model_path = model_path + self.phase = phase + self.is_train_proto = is_train_proto + self.input_shape = input_shape + self.load() + + def load(self): + self.model = get_caffe_resolver().NetParameter() + with open(self.model_path, 'r') as f: + text_format.Merge(f.read(), self.model) + if self.is_train_proto: + self.process_train_proto() + + def process_train_proto(self): + layers = self.model.layer or self.model.layers + delete_layer = set() + split_op_map = dict() + loss_layers = [layer for layer in layers if NodeKind.map_raw_kind(layer.type) in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss)] + a = [layers.remove(layer) for layer in layers[:] if layer in loss_layers[:-1] or NodeKind.map_raw_kind(layer.type) in LAYER_IN_TRAIN_PROTO] + for layer in layers[:]: + if 'label' in layer.bottom: + if NodeKind.map_raw_kind(layer.type) in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss): + continue + elif NodeKind.map_raw_kind(layer.type) == NodeKind.Split: + for item in layer.top: + delete_layer.add(item) + layers.remove(layer) + elif NodeKind.map_raw_kind(layer.type) == NodeKind.Split: + for item in layer.top: + split_op_map[item] = layer.bottom[0] + layers.remove(layer) + + for layer in layers[:]: + for item in delete_layer: + if item in layer.bottom: + layers.remove(layer) + break + for key, value in split_op_map.items(): + if key in layer.bottom: + layer.bottom.remove(key) + layer.bottom.append(value) + self.model.input.append('data') + self.model.input_dim.extend(self.input_shape) + last_layer = layers[-1] + kind = NodeKind.map_raw_kind(last_layer.type) + if kind in (NodeKind.SoftmaxWithLoss, NodeKind.SigmoidCrossEntropyLoss): + pred = layers.add() + pred.name = 'prob' + pred.top.append('prob') + pred.bottom.append(last_layer.bottom[0]) + if kind == NodeKind.SoftmaxWithLoss: + pred.type = NodeKind.Softmax if self.model.layer else 20 # competiable with old version caffe proto + elif kind == NodeKind.SigmoidCrossEntropyLoss: + pred.type = NodeKind.Sigmoid if self.model.layer else 19 + layers.remove(last_layer) + + def filter_layers(self, layers): + phase_map = {0: 'train', 1: 'test'} + filtered_layer_names = set() + filtered_layers = [] + for layer in layers: + phase = self.phase + 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 != self.phase) + # 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 == LayerType.Dropout) + if (not exclude): + exclude = (layer.type == LayerType.Silence) + if not exclude: + if layer.name in filtered_layer_names: + for i in range(1, len(filtered_layer_names)): + new_name = layer.name + '_%s' % i + if new_name not in filtered_layer_names: + layer.name = new_name + break + filtered_layer_names.add(layer.name) + filtered_layers.append(layer) + return filtered_layers + + def make_node(self, layer): + kind = NodeKind.map_raw_kind(layer.type) + if kind is None: + # TODO: raise error + pass + node = CaffeNode(layer.name, kind, layer=layer) + node.output.append(layer.name.replace('/', '_')) + node.output.extend(layer.top[1:]) + return node + + def make_input_node(self): + nodes = [CaffeNode(name, NodeKind.Data) for name in self.model.input] + if len(nodes): + input_dim = list(map(int, self.model.input_dim)) + if not input_dim: + if len(self.model.input_shape) > 0: + input_dim = list(map(int, self.model.input_shape[0].dim)) + else: + # TODO: raise error + pass + for node in nodes: + node.output_shape = tuple(input_dim) + node.output.append('data') + return nodes + + def build(self): + layers = self.model.layers or self.model.layer + layers = self.filter_layers(layers) + nodes = self.make_input_node() + nodes += [self.make_node(layer) for layer in layers] + graph = CaffeGraph(nodes=nodes, name=self.model.name) + node_outputs = {} + for idx, layer in enumerate(layers): + node = graph.get_node(layer.name) + for input_name in layer.bottom: + assert input_name != layer.name + parent_node = node_outputs.get(input_name) + if (parent_node is None) or (parent_node==node): + parent_node = graph.get_node(input_name) + if parent_node.layer: + for i, output in enumerate(parent_node.layer.top): + if input_name == output: + node.add_parent(parent_node, i) + else: + node.add_parent(parent_node, 0) + for output_name in layer.top: + if output_name == layer.name: + continue + node_outputs[output_name] = node + graph.compute_output_shapes(self.model) + return graph