modifed to issue #85

microsoft · May 3, 2018 · b8adc23 · b8adc23
1 parent 0ee5434
commit b8adc23
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 8 deletions.
diff --git a/mmdnn/conversion/mxnet/mxnet_parser.py b/mmdnn/conversion/mxnet/mxnet_parser.py
@@ -898,6 +898,9 @@ def _convert_scalar_operator(self, source_node, new_op):
         value = source_node.get_attr('scalar')
         value_node = self.IR_graph.node.add()
         value_node.name = source_node.real_name + "_second"
+        # left strip the "_" at the beginning of the name
+        # Issue #85, #135
+        value_node.name = value_node.name.lstrip('_')
         value_node.op = 'Constant'
         self.set_weight(value_node.name, 'value', np.array([value], np.float32))
 

diff --git a/mmdnn/conversion/tensorflow/tensorflow_emitter.py b/mmdnn/conversion/tensorflow/tensorflow_emitter.py
@@ -240,14 +240,33 @@ def emit_FullyConnected(self, IR_node):
             bias_str = "bias_initializer = tf.constant_initializer(__weights_dict['{}']['bias']), ".format(IR_node.name)
         else: bias_str = ""
 
-        code = "{:<15} = tf.layers.dense({}, {}, {}{}use_bias = {})".format(
-            IR_node.variable_name,
-            self.parent_variable_name(IR_node),
-            IR_node.layer.attr['units'].i,
-            kernel_str,
-            bias_str,
-            IR_node.layer.attr['use_bias'].b)
-        self.add_body(1, code)
+        # check whether flatten operator should be added
+        parent = self.IR_graph.get_parent(IR_node.name, [0])
+        parent_shape = shape_to_list(parent.get_attr('_output_shapes')[0])
+        if len(parent_shape) > 2:
+            # flatten is needed
+            self.add_body(1, "{:<15} = tf.contrib.layers.flatten({})".format(
+                IR_node.variable_name + '_flatten',
+                self.parent_variable_name(IR_node)))
+
+            code = "{:<15} = tf.layers.dense({}, {}, {}{}use_bias = {})".format(
+                IR_node.variable_name,
+                IR_node.variable_name + '_flatten',
+                IR_node.layer.attr['units'].i,
+                kernel_str,
+                bias_str,
+                IR_node.layer.attr['use_bias'].b)
+            self.add_body(1, code)
+
+        else:
+            code = "{:<15} = tf.layers.dense({}, {}, {}{}use_bias = {})".format(
+                IR_node.variable_name,
+                self.parent_variable_name(IR_node),
+                IR_node.layer.attr['units'].i,
+                kernel_str,
+                bias_str,
+                IR_node.layer.attr['use_bias'].b)
+            self.add_body(1, code)
 
 
     def emit_Flatten(self, IR_node):