Migrate to Python 3, make V1 compatible.

convert.py

@@ -31,7 +31,7 @@ def convert(def_path, caffemodel_path, data_output_path, code_output_path, phase
                 np.save(data_out, data)
         if code_output_path:
             print_stderr('Saving source...')
-            with open(code_output_path, 'wb') as src_out:
+            with open(code_output_path, 'w') as src_out:
                 src_out.write(transformer.transform_source())
         print_stderr('Done.')
     except KaffeError as err:

kaffe/caffe/resolver.py

@@ -8,19 +8,10 @@ def __init__(self):
 
     def import_caffe(self):
         self.caffe = None
-        try:
-            # Try to import PyCaffe first
-            import caffe
-            self.caffe = caffe
-        except ImportError:
-            # Fall back to the protobuf implementation
-            from . import caffepb
-            self.caffepb = caffepb
-            show_fallback_warning()
-        if self.caffe:
-            # Use the protobuf code from the imported distribution.
-            # This way, Caffe variants with custom layers will work.
-            self.caffepb = self.caffe.proto.caffe_pb2
+        # Fall back to the protobuf implementation
+        from . import caffepb
+        self.caffepb = caffepb
+        show_fallback_warning()
         self.NetParameter = self.caffepb.NetParameter
 
     def has_pycaffe(self):

kaffe/core.py

@@ -278,7 +278,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)
         if kind is None:
             raise KaffeError('Unknown layer type encountered: %s'%layer.type)
         return Node(layer.name, kind, layer=layer)

kaffe/graph.py

@@ -1,6 +1,6 @@
 from google.protobuf import text_format
 
-from .caffe import get_caffe_resolver
+from .caffe import get_caffe_resolver, caffepb
 from .errors import KaffeError, print_stderr
 from .layers import LayerAdapter, LayerType, NodeKind, NodeDispatch
 from .shapes import TensorShape
@@ -114,15 +114,16 @@ def __contains__(self, key):
         return key in self.node_lut
 
     def __str__(self):
-        hdr = '{:<20} {:<30} {:>20} {:>20}'.format('Type', 'Name', 'Param', 'Output')
+        hdr = '{:<20} {:<30} {:>20} {:>20}'.format(
+            'Type', 'Name', 'Param', 'Output')
         s = [hdr, '-' * 94]
         for node in self.topologically_sorted():
             # If the node has learned parameters, display the first one's shape.
             # In case of convolutions, this corresponds to the weights.
-            data_shape = node.data[0].shape if node.data else '--'
-            out_shape = node.output_shape or '--'
-            s.append('{:<20} {:<30} {:>20} {:>20}'.format(node.kind, node.name, data_shape,
-                                                          tuple(out_shape)))
+            data_shape = str(node.data[0].shape) if node.data else '--'
+            out_shape = str(tuple(node.output_shape)) or '--'
+            s.append('{:<20} {:<30} {:>20} {:>20}'.format(
+                node.kind, node.name, data_shape, out_shape))
         return '\n'.join(s)
 
 
@@ -142,7 +143,7 @@ def __init__(self, def_path, phase='test'):
     def load(self):
         '''Load the layer definitions from the prototxt.'''
         self.params = get_caffe_resolver().NetParameter()
-        with open(self.def_path, 'rb') as def_file:
+        with open(self.def_path, 'r') as def_file:
             text_format.Merge(def_file.read(), self.params)
 
     def filter_layers(self, layers):
@@ -171,7 +172,7 @@ def filter_layers(self, layers):
 
     def make_node(self, layer):
         '''Create a graph node for the given layer.'''
-        kind = NodeKind.map_raw_kind(layer.type)
+        kind = NodeKind.map_raw_kind(layer)
         if kind is None:
             raise KaffeError('Unknown layer type encountered: %s' % layer.type)
         # We want to use the layer's top names (the "output" names), rather than the

kaffe/layers.py

@@ -1,6 +1,7 @@
 import re
 import numbers
 from collections import namedtuple
+from .caffe import caffepb
 
 from .shapes import *
 
@@ -51,14 +52,59 @@
     'Threshold': shape_identity,
 }
 
+V1_TO_NEW = {
+    35: 'AbsVal',
+    1: 'Accuracy',
+    30: 'ArgMax',
+    2: 'BNLL',
+    3: 'Concat',
+    37: 'ContrastiveLoss',
+    4: 'Convolution',
+    5: 'Data',
+    39: 'Deconvolution',
+    6: 'Dropout',
+    32: 'DummyData',
+    7: 'EuclideanLoss',
+    25: 'Eltwise',
+    38: 'Exp',
+    8: 'Flatten',
+    9: 'HDF5Data',
+    10: 'HDF5Output',
+    28: 'HingeLoss',
+    11: 'Im2col',
+    12: 'ImageData',
+    13: 'InfogainLoss',
+    14: 'InnerProduct',
+    15: 'LRN',
+    29: 'MemoryData',
+    16: 'MultinomialLogisticLoss',
+    34: 'MVN',
+    17: 'Pooling',
+    26: 'Power',
+    18: 'ReLU',
+    19: 'Sigmoid',
+    27: 'SigmoidCrossEntropyLoss',
+    36: 'Silence',
+    20: 'Softmax',
+    21: 'SoftmaxLoss',
+    22: 'Split',
+    33: 'Slice',
+    23: 'TanH',
+    24: 'WindowData',
+    31: 'Threshold',
+}
+
 LAYER_TYPES = list(LAYER_DESCRIPTORS.keys())
 
 LayerType = type('LayerType', (), {t: t for t in LAYER_TYPES})
 
 class NodeKind(LayerType):
 
     @staticmethod
-    def map_raw_kind(kind):
+    def map_raw_kind(layer):
+        kind = layer.type
+        if isinstance(layer, caffepb.V1LayerParameter):
+            kind = V1_TO_NEW[layer.type]
         if kind in LAYER_TYPES:
             return kind
         return None

kaffe/tensorflow/transformer.py

@@ -44,7 +44,7 @@ def __init__(self, op, *args, **kwargs):
 
     def format(self, arg):
         '''Returns a string representation for the given value.'''
-        return "'%s'" % arg if isinstance(arg, basestring) else str(arg)
+        return "'%s'" % arg if isinstance(arg, str) else str(arg)
 
     def pair(self, key, value):
         '''Returns key=formatted(value).'''
@@ -53,7 +53,7 @@ def pair(self, key, value):
     def emit(self):
         '''Emits the Python source for this node.'''
         # Format positional arguments
-        args = map(self.format, self.args)
+        args = list(map(self.format, self.args))
         # Format any keyword arguments
         if self.kwargs:
             args += [self.pair(k, v) for k, v in self.kwargs]