caffe global pooling supported. caffe xception -> tf tested.

Author: kitstar

mmdnn/conversion/caffe/graph.py

@@ -105,7 +105,7 @@
 
 LayerType = type('LayerType', (), {t : t for t in LAYER_TYPES})
 
-KernelParameters = namedtuple('KernelParameters', ['k_h', 'k_w', 's_h', 's_w', 'p_h', 'p_w'])
+KernelParameters = namedtuple('KernelParameters', ['global_pooling', 'k_h', 'k_w', 's_h', 's_w', 'p_h', 'p_w'])
 
 class NodeKind(LayerType):
 
@@ -185,45 +185,50 @@ def get_kernel_value(scalar, repeated, idx, default=None):
     def kernel_parameters(self):
         assert self.kind in (NodeKind.Convolution, NodeKind.Pooling, NodeKind.Deconvolution)
         params = self.parameters
-        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)
+        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_h, params.pad, 1, default=0)
-        return KernelParameters(k_h, k_w, s_h, s_w, p_h, p_w)
+        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 = []
@@ -263,11 +268,11 @@ def compute_output_shapes(self, model):
         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:
@@ -316,7 +321,7 @@ def load(self):
             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()
@@ -359,7 +364,7 @@ def process_train_proto(self):
             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()
@@ -388,7 +393,7 @@ def filter_layers(self, layers):
         return filtered_layers
 
     def make_node(self, layer):
-        kind = NodeKind.map_raw_kind(layer.type)            
+        kind = NodeKind.map_raw_kind(layer.type)
         if kind is None:
             # TODO: raise error
             pass

mmdnn/conversion/caffe/mapper.py

@@ -36,14 +36,27 @@ def _convert_output_shape(cls, kwargs, node):
     def get_kernel_params(cls, node, input_shape):
         kwargs = {}
 
-        o_h_caffe = node.output_shape.height
-        o_h_tf = (input_shape.height + node.kernel_parameters.p_h * 2 - node.kernel_parameters.k_h + 1) // node.kernel_parameters.s_h
+        if node.kernel_parameters.global_pooling:
+            kwargs['kernel_shape'] = [1, input_shape.height, input_shape.width, 1]
+            kwargs['pads'] = [0] * 8
 
-        o_w_caffe = node.output_shape.width
-        o_w_tf = (input_shape.width + node.kernel_parameters.p_w * 2 - node.kernel_parameters.k_w + 1) // node.kernel_parameters.s_w
+        else:
+            from mmdnn.conversion.caffe.graph import NodeKind
+            if node.kind == NodeKind.Pooling:
+                kwargs['kernel_shape'] = [1, node.kernel_parameters.k_h, node.kernel_parameters.k_w, 1]
+            elif node.kind == NodeKind.Convolution:
+                pass
+            else:
+                raise ValueError
+
+            o_h_caffe = node.output_shape.height
+            o_h_tf = (input_shape.height + node.kernel_parameters.p_h * 2 - node.kernel_parameters.k_h + 1) // node.kernel_parameters.s_h
+            o_w_caffe = node.output_shape.width
+            o_w_tf = (input_shape.width + node.kernel_parameters.p_w * 2 - node.kernel_parameters.k_w + 1) // node.kernel_parameters.s_w
+
+            kwargs['pads'] = [0, node.kernel_parameters.p_h, node.kernel_parameters.p_w, 0] + \
+                    [0, node.kernel_parameters.p_h + o_h_caffe - o_h_tf, node.kernel_parameters.p_w + o_w_caffe - o_w_tf, 0]
 
-        kwargs['pads'] = [0, node.kernel_parameters.p_h, node.kernel_parameters.p_w, 0] + \
-                  [0, node.kernel_parameters.p_h + o_h_caffe - o_h_tf, node.kernel_parameters.p_w + o_w_caffe - o_w_tf, 0]
         kwargs['strides'] = [1, node.kernel_parameters.s_h, node.kernel_parameters.s_w, 1]
         cls._convert_output_shape(kwargs, node)
 
@@ -117,7 +130,6 @@ def map_pooling(cls, node):
         else:
             # Stochastic pooling, for instance.
             raise ConversionError('Unsupported pooling type.')
-        kwargs['kernel_shape'] = [1, node.kernel_parameters.k_h, node.kernel_parameters.k_w, 1]
         cls._convert_output_shape(kwargs, node)
         return Node.create('Pool', **kwargs)