Add option to save standalone model

README.md

@@ -13,6 +13,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

@@ -7,6 +7,10 @@
 from kaffe import KaffeError, print_stderr
 from kaffe.tensorflow import TensorFlowTransformer
 
+import shutil
+import tensorflow as tf
+from tensorflow.python.tools.freeze_graph import freeze_graph
+
 
 def fatal_error(msg):
     print_stderr(msg)
@@ -16,25 +20,90 @@ 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):
     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:
+        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')
+            os.makedirs(temp_folder)
+
+            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 data_out:
+                    np.save(data_out, data)
+
+            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.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.train.Saver()
+            saver.save(sess, checkpoint_path)
+
+            print_stderr('Saving graph definition as protobuf...')
+            tf.train.write_graph(sess.graph.as_graph_def(), graph_folder, graph_name, False)
+
+            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...')
+            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, '')
+
+            shutil.rmtree(temp_folder)
+
         print_stderr('Done.')
     except KaffeError as err:
         fatal_error('Error encountered: {}'.format(err))
@@ -46,14 +115,15 @@ 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')
     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)
 
 
 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/tensorflow/network.py

@@ -212,7 +212,7 @@ def softmax(self, input, name):
                 input = tf.squeeze(input, squeeze_dims=[1, 2])
             else:
                 raise ValueError('Rank 2 tensor input expected for softmax!')
-        return tf.nn.softmax(input, name)
+        return tf.nn.softmax(input, name=name)
 
     @layer
     def batch_normalization(self, input, name, scale_offset=True, relu=False):

kaffe/tensorflow/transformer.py

@@ -281,5 +281,7 @@ def transform_source(self):
             mapper = TensorFlowMapper(self.graph)
             chains = mapper.map()
             emitter = TensorFlowEmitter()
+            if not self.graph.name:
+                self.graph.name = 'MyNet'
             self.source = emitter.emit(self.graph.name, chains)
         return self.source