Support PReLU Layer

.gitignore

@@ -8,3 +8,12 @@
 # Python cache
 *.pyc
 
+# Ignore Results
+*.pb
+*.pbtxt
+*.mat
+.tmp/
+
+# Ignore Development setting
+.vscode/
+.ipynb_checkpoints/

README.md

@@ -2,6 +2,24 @@
 
 Convert [Caffe](https://github.com/BVLC/caffe/) models to [TensorFlow](https://github.com/tensorflow/tensorflow).
 
+## Usage with standalone model
+
+### 1 - Install caffe-tensorflow
+    git clone https://github.com/linkfluence/caffe-tensorflow
+    # Optional: create a Python 2.7 env and activate it
+    # This fork has only be tested with Python 2.7
+
+### 2 - (Optional) Switch to Tensorflow CPU
+You might bump into memory issues if you don't have enough memory. In this case just uninstall `tensorflow-gpu` and install `tensorflow`
+
+### 3 - Convert your model
+    python convert.py --caffemodel ./model.caffemodel ./model.prototxt --data-output-path ./output.mat --code-output-path ./output.py --standalone-output-path ./standalone.pb`
+
+### 4 - (Optional) Re-install Tensorflow GPU 
+
+### 5- Use the standalone.pb file 
+It contains the weights and the architecture of the network.
+
 ## Usage
 
 Run `convert.py` to convert an existing Caffe model to TensorFlow.
@@ -13,6 +31,8 @@ The output consists of two files:
 1. A data file (in NumPy's native format) containing the model's learned parameters.
 2. A Python class that constructs the model's graph.
 
+Alternatively, you can save a standalone GraphDef model file containing the model's graph and learned parameters. 
+
 ### Examples
 
 See the [examples](examples/) folder for more details.

convert.py

@@ -2,8 +2,14 @@
 
 import os
 import sys
-import numpy as np
 import argparse
+
+import shutil
+import pickle
+import tensorflow as tf
+from tensorflow.python.tools.freeze_graph import freeze_graph
+from tensorflow.python.tools import optimize_for_inference_lib
+
 from kaffe import KaffeError, print_stderr
 from kaffe.tensorflow import TensorFlowTransformer
 
@@ -16,25 +22,113 @@ def fatal_error(msg):
 def validate_arguments(args):
     if (args.data_output_path is not None) and (args.caffemodel is None):
         fatal_error('No input data path provided.')
-    if (args.caffemodel is not None) and (args.data_output_path is None):
+    if (args.caffemodel is not None) and (args.data_output_path is None) and \
+        (args.standalone_output_path is None):
         fatal_error('No output data path provided.')
-    if (args.code_output_path is None) and (args.data_output_path is None):
+    if (args.code_output_path is None) and (args.data_output_path is None) and \
+        (args.standalone_output_path is None):
         fatal_error('No output path specified.')
 
 
-def convert(def_path, caffemodel_path, data_output_path, code_output_path, phase):
+def convert(def_path, caffemodel_path, data_output_path, code_output_path, standalone_output_path,
+            phase, freeze):
     try:
+        sess = tf.InteractiveSession()
         transformer = TensorFlowTransformer(def_path, caffemodel_path, phase=phase)
         print_stderr('Converting data...')
-        if caffemodel_path is not None:
+        if data_output_path is not None:
             data = transformer.transform_data()
             print_stderr('Saving data...')
-            with open(data_output_path, 'wb') as data_out:
-                np.save(data_out, data)
-        if code_output_path:
+            with open(data_output_path, 'wb') as handle:
+                pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)
+        if code_output_path is not None:
             print_stderr('Saving source...')
             with open(code_output_path, 'wb') as src_out:
                 src_out.write(transformer.transform_source())
+
+        if standalone_output_path:
+            filename, _ = os.path.splitext(os.path.basename(standalone_output_path))
+            temp_folder = os.path.join(os.path.dirname(standalone_output_path), '.tmp')
+            if not os.path.exists(temp_folder):
+                os.makedirs(temp_folder)
+                shutil.rmtree(temp_folder) # Delete old graphs
+
+            if data_output_path is None:
+                data = transformer.transform_data()
+                print_stderr('Saving data...')
+                data_output_path = os.path.join(temp_folder, filename) + '.npy'
+                with open(data_output_path, 'wb') as handle:
+                    pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)
+
+            if code_output_path is None:
+                print_stderr('Saving source...')
+                code_output_path = os.path.join(temp_folder, filename) + '.py'
+                with open(code_output_path, 'wb') as src_out:
+                    src_out.write(transformer.transform_source())
+
+            checkpoint_path = os.path.join(temp_folder, filename + '.ckpt')
+            graph_name = os.path.basename(standalone_output_path)
+            graph_folder = os.path.dirname(standalone_output_path)
+            input_node = transformer.graph.nodes[0].name
+            output_node = transformer.graph.nodes[-1].name
+            tensor_shape = transformer.graph.get_node(input_node).output_shape
+            tensor_shape_list = [tensor_shape.batch_size, tensor_shape.height,
+                                 tensor_shape.width, tensor_shape.channels]
+
+            sys.path.append(os.path.dirname(code_output_path))
+            module = os.path.splitext(os.path.basename(code_output_path))[0]
+            class_name = transformer.graph.name
+            KaffeNet = getattr(__import__(module), class_name)
+
+            data_placeholder = tf.compat.v1.placeholder(
+                tf.float32, tensor_shape_list, name=input_node)
+            net = KaffeNet({input_node: data_placeholder})
+
+            # load weights stored in numpy format
+            net.load(data_output_path, sess)
+
+            print_stderr('Saving checkpoint...')
+            saver = tf.compat.v1.train.Saver()
+            saver.save(sess, checkpoint_path)
+
+            print_stderr('Saving graph definition as protobuf...')
+            tf.io.write_graph(sess.graph.as_graph_def(), graph_folder, graph_name, False)
+            writer = tf.compat.v1.summary.FileWriter('.tmp', sess.graph)
+            writer.close()
+
+            input_graph_path = standalone_output_path
+            input_saver_def_path = ""
+            input_binary = True
+            input_checkpoint_path = checkpoint_path
+            output_node_names = output_node
+            restore_op_name = 'save/restore_all'
+            filename_tensor_name = 'save/Const:0'
+            output_graph_path = standalone_output_path
+            clear_devices = True
+
+            print_stderr('Saving standalone model...')
+            output_node_names = '{0}/{0}'.format(output_node_names)
+            if freeze == 'freeze_graph':
+                freeze_graph(input_graph_path, input_saver_def_path,
+                             input_binary, input_checkpoint_path,
+                             output_node_names, restore_op_name,
+                             filename_tensor_name, output_graph_path,
+                             clear_devices, '')
+            elif freeze == 'optimize_for_inference':
+                graph_def = sess.graph.as_graph_def()
+                graph_def = tf.graph_util.convert_variables_to_constants(
+                    sess, graph_def, [output_node_names])
+                graph_def_f32 = optimize_for_inference_lib.optimize_for_inference(
+                    graph_def, ['data'], [output_node_names], tf.float32.as_datatype_enum)
+                tf.train.write_graph(
+                    graph_def_f32, "", standalone_output_path.rsplit('.',1)[0] + '.pb', as_text=False)
+                tf.train.write_graph(
+                    graph_def_f32, "", standalone_output_path.rsplit('.',1)[0] + '.pbtxt', as_text=True)
+
+            #f = shutil.rmtree(temp_folder)
+            writer = tf.compat.v1.summary.FileWriter('.tmp', sess.graph)
+            writer.close()
+
         print_stderr('Done.')
     except KaffeError as err:
         fatal_error('Error encountered: {}'.format(err))
@@ -46,14 +140,21 @@ def main():
     parser.add_argument('--caffemodel', help='Model data (.caffemodel) path')
     parser.add_argument('--data-output-path', help='Converted data output path')
     parser.add_argument('--code-output-path', help='Save generated source to this path')
+    parser.add_argument('--standalone-output-path',
+                        help='Save generated standalone tensorflow model to this path')
     parser.add_argument('-p',
                         '--phase',
                         default='test',
                         help='The phase to convert: test (default) or train')
+    parser.add_argument('-fz',
+                        '--freeze',
+                        default=None,
+                        help="""Freeze option for inference: No (default),
+                                freeze_graph or optimize_for_inference(e.g. for OpenCV)""")
     args = parser.parse_args()
     validate_arguments(args)
     convert(args.def_path, args.caffemodel, args.data_output_path, args.code_output_path,
-            args.phase)
+            args.standalone_output_path, args.phase, args.freeze)
 
 
 if __name__ == '__main__':

examples/mnist/README.md

@@ -34,3 +34,11 @@ with tf.Session() as sesh:
     # Forward pass
     output = sesh.run(net.get_output(), ...)
 ```
+
+#### Standalone model file:
+
+You can save a standalone GraphDef model file as follows:
+
+    $ ./convert.py examples/mnist/lenet.prototxt --caffemodel examples/mnist/lenet_iter_10000.caffemodel --standalone-output-path=mynet.pb
+
+This generates a protobuf file named `mynet.pb` containing the model's graph and parameters. The [TensorFlow Image Recognition tutorial](https://www.tensorflow.org/versions/r0.11/tutorials/image_recognition/index.html) shows how to use models constructed in this way in [Python](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/models/image/imagenet) or [C++](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/label_image).

kaffe/caffe/resolver.py

@@ -14,8 +14,8 @@ def import_caffe(self):
             self.caffe = caffe
         except ImportError:
             # Fall back to the protobuf implementation
-            from . import caffepb
-            self.caffepb = caffepb
+            from . import caffe_pb2
+            self.caffepb = caffe_pb2
             show_fallback_warning()
         if self.caffe:
             # Use the protobuf code from the imported distribution.

kaffe/graph.py

@@ -119,7 +119,12 @@ def __str__(self):
         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 '--'
+            if node.data is None:
+                data_shape = '--'
+            elif isinstance(node.data, dict):
+                data_shape = node.data['weights'].shape#'dict({})'.format(node.data.keys())
+            else:
+                data_shape = node.data[0].shape
             out_shape = node.output_shape or '--'
             s.append('{:<20} {:<30} {:>20} {:>20}'.format(node.kind, node.name, data_shape,
                                                           tuple(out_shape)))

kaffe/layers.py

@@ -38,6 +38,7 @@
     'Pooling': shape_pool,
     'Power': shape_identity,
     'ReLU': shape_identity,
+    'PReLU': shape_identity,
     'Scale': shape_identity,
     'Sigmoid': shape_identity,
     'SigmoidCrossEntropyLoss': shape_scalar,
@@ -81,7 +82,7 @@ class NodeDispatch(object):
 
     @staticmethod
     def get_handler_name(node_kind):
-        if len(node_kind) <= 4:
+        if len(node_kind) <= 4 or node_kind == 'PReLU':
             # A catch-all for things like ReLU and tanh
             return node_kind.lower()
         # Convert from CamelCase to under_scored