Enable python style check

.gitlab-ci.yml

 stages:
   - cpplint
+  - pycodestyle
   - ops_test
   - ops_benchmark
 
@@ -7,7 +8,12 @@ cpplint:
   stage: cpplint
   script:
     - curl -o cpplint.py https://raw.githubusercontent.com/google/styleguide/gh-pages/cpplint/cpplint.py
-    - python cpplint.py --linelength=80 --counting=detailed $(find mace -name *.h -or -name *.cc)
+    - python cpplint.py --linelength=80 --counting=detailed $(find mace -name "*.h" -or -name "*.cc")
+
+pycodestyle:
+  stage: pycodestyle
+  script:
+    - pycodestyle $(find -name "*.py")
 
 ops_test:
   stage: ops_test

docker/Dockerfile

@@ -113,7 +113,8 @@ RUN pip install -i http://pypi.douban.com/simple/ --trusted-host pypi.douban.com
     scipy \
     jinja2 \
     pyyaml \
-    sh 
+    sh \
+    pycodestyle
 
 # Download tensorflow tools
 RUN wget http://cnbj1-inner-fds.api.xiaomi.net/mace/tool/transform_graph && \

mace/python/tools/binary_codegen.py

@@ -16,74 +16,75 @@ FLAGS = None
 
 
 def generate_cpp_source():
-  data_map = {}
-  for binary_dir in FLAGS.binary_dirs.split(","):
-    binary_path = os.path.join(binary_dir, FLAGS.binary_file_name)
-    if not os.path.exists(binary_path):
-      continue
+    data_map = {}
+    for binary_dir in FLAGS.binary_dirs.split(","):
+        binary_path = os.path.join(binary_dir, FLAGS.binary_file_name)
+        if not os.path.exists(binary_path):
+            continue
 
-    with open(binary_path, "rb") as f:
-      binary_array = np.fromfile(f, dtype=np.uint8)
+        with open(binary_path, "rb") as f:
+            binary_array = np.fromfile(f, dtype=np.uint8)
 
-    print "Generate binary from", binary_path
-    idx = 0
-    size, = struct.unpack("Q", binary_array[idx:idx+8])
-    idx += 8
-    for _ in xrange(size):
-      key_size, = struct.unpack("i", binary_array[idx:idx+4])
-      idx += 4
-      key, = struct.unpack(str(key_size) + "s", binary_array[idx:idx+key_size])
-      idx += key_size
-      params_size, = struct.unpack("i", binary_array[idx:idx+4])
-      idx += 4
-      data_map[key] = []
-      count = params_size / 4
-      params = struct.unpack(str(count) + "i", binary_array[idx:idx+params_size])
-      for i in params:
-        data_map[key].append(i)
-      idx += params_size
+        print "Generate binary from", binary_path
+        idx = 0
+        size, = struct.unpack("Q", binary_array[idx:idx + 8])
+        idx += 8
+        for _ in xrange(size):
+            key_size, = struct.unpack("i", binary_array[idx:idx + 4])
+            idx += 4
+            key, = struct.unpack(
+                str(key_size) + "s", binary_array[idx:idx + key_size])
+            idx += key_size
+            params_size, = struct.unpack("i", binary_array[idx:idx + 4])
+            idx += 4
+            data_map[key] = []
+            count = params_size / 4
+            params = struct.unpack(
+                str(count) + "i", binary_array[idx:idx + params_size])
+            for i in params:
+                data_map[key].append(i)
+            idx += params_size
+
+    env = jinja2.Environment(loader=jinja2.FileSystemLoader(sys.path[0]))
+    return env.get_template('str2vec_maps.cc.jinja2').render(
+        maps=data_map,
+        data_type='unsigned int',
+        variable_name=FLAGS.variable_name)
 
-  env = jinja2.Environment(loader=jinja2.FileSystemLoader(sys.path[0]))
-  return env.get_template('str2vec_maps.cc.jinja2').render(
-    maps = data_map,
-    data_type = 'unsigned int',
-    variable_name = FLAGS.variable_name
-  )
 
 def main(unused_args):
-  cpp_binary_source = generate_cpp_source()
-  if os.path.isfile(FLAGS.output_path):
-    os.remove(FLAGS.output_path)
-  w_file = open(FLAGS.output_path, "w")
-  w_file.write(cpp_binary_source)
-  w_file.close()
+    cpp_binary_source = generate_cpp_source()
+    if os.path.isfile(FLAGS.output_path):
+        os.remove(FLAGS.output_path)
+    w_file = open(FLAGS.output_path, "w")
+    w_file.write(cpp_binary_source)
+    w_file.close()
+
 
 def parse_args():
-  """Parses command line arguments."""
-  parser = argparse.ArgumentParser()
-  parser.add_argument(
-      "--binary_dirs",
-      type=str,
-      default="",
-      help="The binaries file path.")
-  parser.add_argument(
-      "--binary_file_name",
-      type=str,
-      default="mace_run.config",
-      help="The binary file name.")
-  parser.add_argument(
-      "--output_path",
-      type=str,
-      default="",
-      help="The path of generated C++ source file which contains the binary.")
-  parser.add_argument(
-    "--variable_name",
-    type=str,
-    default="kTuningParamsData",
-    help="global variable name.")
-  return parser.parse_known_args()
+    """Parses command line arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--binary_dirs", type=str, default="", help="The binaries file path.")
+    parser.add_argument(
+        "--binary_file_name",
+        type=str,
+        default="mace_run.config",
+        help="The binary file name.")
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="",
+        help="The path of generated C++ source file which contains the binary."
+    )
+    parser.add_argument(
+        "--variable_name",
+        type=str,
+        default="kTuningParamsData",
+        help="global variable name.")
+    return parser.parse_known_args()
 
 
 if __name__ == '__main__':
-  FLAGS, unparsed = parse_args()
-  main(unused_args=[sys.argv[0]] + unparsed)
+    FLAGS, unparsed = parse_args()
+    main(unused_args=[sys.argv[0]] + unparsed)

mace/python/tools/convert_util.py

@@ -26,4 +26,3 @@ def tf_dtype_2_mace_dtype(tf_dtype):
     if not mace_dtype:
         raise Exception("Not supported tensorflow dtype: " + tf_dtype)
     return mace_dtype
-

mace/python/tools/converter.py

@@ -4,176 +4,166 @@ import hashlib
 import os.path
 from mace.python.tools import source_converter_lib
 
-# ./bazel-bin/mace/python/tools/tf_converter --model_file quantized_test.pb --output quantized_test_dsp.pb --runtime dsp --input_dim input_node,1,28,28,3
+# ./bazel-bin/mace/python/tools/tf_converter --model_file quantized_test.pb \
+#                                            --output quantized_test_dsp.pb \
+#                                            --runtime dsp \
+#                                            --input_dim input_node,1,28,28,3
 
 FLAGS = None
 
+
 def file_checksum(fname):
-  hash_func = hashlib.sha256()
-  with open(fname, "rb") as f:
-    for chunk in iter(lambda: f.read(4096), b""):
-      hash_func.update(chunk)
-  return hash_func.hexdigest()
+    hash_func = hashlib.sha256()
+    with open(fname, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_func.update(chunk)
+    return hash_func.hexdigest()
+
 
 def main(unused_args):
-  if not os.path.isfile(FLAGS.model_file):
-    print("Input graph file '" + FLAGS.model_file + "' does not exist!")
-    sys.exit(-1)
-
-  model_checksum = file_checksum(FLAGS.model_file)
-  if FLAGS.model_checksum != "" and FLAGS.model_checksum != model_checksum:
-    print("Model checksum mismatch: %s != %s" % (model_checksum, FLAGS.model_checksum))
-    sys.exit(-1)
-
-  if FLAGS.platform == 'caffe':
-    if not os.path.isfile(FLAGS.weight_file):
-      print("Input weight file '" + FLAGS.weight_file + "' does not exist!")
-      sys.exit(-1)
-
-    weight_checksum = file_checksum(FLAGS.weight_file)
-    if FLAGS.weight_checksum != "" and FLAGS.weight_checksum != weight_checksum:
-      print("Weight checksum mismatch: %s != %s" % (weight_checksum, FLAGS.weight_checksum))
-      sys.exit(-1)
-
-    if FLAGS.runtime == 'dsp':
-      print("DSP not support caffe model yet.")
-      sys.exit(-1)
-
-    from mace.python.tools import caffe_converter_lib
-    output_graph_def = caffe_converter_lib.convert_to_mace_pb(
-      FLAGS.model_file, FLAGS.weight_file, FLAGS.input_node, FLAGS.input_shape, FLAGS.output_node,
-      FLAGS.data_type, FLAGS.runtime, FLAGS.winograd)
-  elif FLAGS.platform == 'tensorflow':
-    if FLAGS.runtime == 'dsp':
-      from mace.python.tools import tf_dsp_converter_lib
-      output_graph_def = tf_dsp_converter_lib.convert_to_mace_pb(
-        FLAGS.model_file, FLAGS.input_node, FLAGS.output_node, FLAGS.dsp_mode)
+    if not os.path.isfile(FLAGS.model_file):
+        print("Input graph file '" + FLAGS.model_file + "' does not exist!")
+        sys.exit(-1)
+
+    model_checksum = file_checksum(FLAGS.model_file)
+    if FLAGS.model_checksum != "" and FLAGS.model_checksum != model_checksum:
+        print("Model checksum mismatch: %s != %s" % (model_checksum,
+                                                     FLAGS.model_checksum))
+        sys.exit(-1)
+
+    if FLAGS.platform == 'caffe':
+        if not os.path.isfile(FLAGS.weight_file):
+            print("Input weight file '" + FLAGS.weight_file +
+                  "' does not exist!")
+            sys.exit(-1)
+
+        weight_checksum = file_checksum(FLAGS.weight_file)
+        if FLAGS.weight_checksum != "" and \
+                FLAGS.weight_checksum != weight_checksum:
+            print("Weight checksum mismatch: %s != %s" %
+                  (weight_checksum, FLAGS.weight_checksum))
+            sys.exit(-1)
+
+        if FLAGS.runtime == 'dsp':
+            print("DSP not support caffe model yet.")
+            sys.exit(-1)
+
+        from mace.python.tools import caffe_converter_lib
+        output_graph_def = caffe_converter_lib.convert_to_mace_pb(
+            FLAGS.model_file, FLAGS.weight_file, FLAGS.input_node,
+            FLAGS.input_shape, FLAGS.output_node, FLAGS.data_type,
+            FLAGS.runtime, FLAGS.winograd)
+    elif FLAGS.platform == 'tensorflow':
+        if FLAGS.runtime == 'dsp':
+            from mace.python.tools import tf_dsp_converter_lib
+            output_graph_def = tf_dsp_converter_lib.convert_to_mace_pb(
+                FLAGS.model_file, FLAGS.input_node, FLAGS.output_node,
+                FLAGS.dsp_mode)
+        else:
+            from mace.python.tools import tf_converter_lib
+            output_graph_def = tf_converter_lib.convert_to_mace_pb(
+                FLAGS.model_file, FLAGS.input_node, FLAGS.input_shape,
+                FLAGS.output_node, FLAGS.data_type, FLAGS.runtime,
+                FLAGS.winograd)
+
+    if FLAGS.output_type == 'source':
+        source_converter_lib.convert_to_source(
+            output_graph_def, model_checksum, FLAGS.template, FLAGS.obfuscate,
+            FLAGS.model_tag, FLAGS.output, FLAGS.runtime,
+            FLAGS.embed_model_data)
     else:
-      from mace.python.tools import tf_converter_lib
-      output_graph_def = tf_converter_lib.convert_to_mace_pb(
-        FLAGS.model_file, FLAGS.input_node, FLAGS.input_shape, FLAGS.output_node,
-        FLAGS.data_type, FLAGS.runtime, FLAGS.winograd)
-
-  if FLAGS.output_type == 'source':
-    source_converter_lib.convert_to_source(output_graph_def, model_checksum, FLAGS.template, FLAGS.obfuscate,
-      FLAGS.model_tag, FLAGS.output, FLAGS.runtime, FLAGS.embed_model_data)
-  else:
-    with open(FLAGS.output, "wb") as f:
-      f.write(output_graph_def.SerializeToString())
-    with open(FLAGS.output + '_txt', "wb") as f:
-      # output_graph_def.ClearField('tensors')
-      f.write(str(output_graph_def))
-  print("Model conversion is completed.")
+        with open(FLAGS.output, "wb") as f:
+            f.write(output_graph_def.SerializeToString())
+        with open(FLAGS.output + '_txt', "wb") as f:
+            # output_graph_def.ClearField('tensors')
+            f.write(str(output_graph_def))
+    print("Model conversion is completed.")
+
 
 def str2bool(v):
-  if v.lower() in ('yes', 'true', 't', 'y', '1'):
-    return True
-  elif v.lower() in ('no', 'false', 'f', 'n', '0'):
-    return False
-  else:
-    raise argparse.ArgumentTypeError('Boolean value expected.')
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
 
 def parse_args():
-  """Parses command line arguments."""
-  parser = argparse.ArgumentParser()
-  parser.register("type", "bool", lambda v: v.lower() == "true")
-  parser.add_argument(
-    "--model_file",
-    type=str,
-    default="",
-    help="TensorFlow \'GraphDef\' file to load, Caffe prototxt file to load.")
-  parser.add_argument(
-    "--weight_file",
-    type=str,
-    default="",
-    help="Caffe data file to load.")
-  parser.add_argument(
-    "--model_checksum",
-    type=str,
-    default="",
-    help="Model file sha256 checksum")
-  parser.add_argument(
-    "--weight_checksum",
-    type=str,
-    default="",
-    help="Weight file sha256 checksum")
-  parser.add_argument(
-    "--output",
-    type=str,
-    default="",
-    help="File to save the output graph to.")
-  parser.add_argument(
-    "--runtime",
-    type=str,
-    default="cpu",
-    help="Runtime: cpu/gpu/dsp")
-  parser.add_argument(
-    "--input_node",
-    type=str,
-    default="input_node",
-    help="e.g., input_node")
-  parser.add_argument(
-    "--output_node",
-    type=str,
-    default="softmax",
-    help="e.g., softmax")
-  parser.add_argument(
-    "--data_type",
-    type=str,
-    default='DT_FLOAT',
-    help="e.g., DT_HALF/DT_FLOAT")
-  parser.add_argument(
-    "--output_type",
-    type=str,
-    default="pb",
-    help="output type: source/pb")
-  parser.add_argument(
-    "--template",
-    type=str,
-    default="",
-    help="template path")
-  parser.add_argument(
-    "--obfuscate",
-    type=str2bool,
-    nargs='?',
-    const=False,
-    default=False,
-    help="obfuscate model names")
-  parser.add_argument(
-    "--model_tag",
-    type=str,
-    default="",
-    help="model tag for generated function and namespace")
-  parser.add_argument(
-    "--winograd",
-    type=str2bool,
-    nargs='?',
-    const=False,
-    default=False,
-    help="open winograd convolution or not")
-  parser.add_argument(
-    "--dsp_mode",
-    type=int,
-    default=0,
-    help="dsp run mode, defalut=0")
-  parser.add_argument(
-    "--input_shape",
-    type=str,
-    default="",
-    help="input shape.")
-  parser.add_argument(
-    "--platform",
-    type=str,
-    default="tensorflow",
-    help="tensorflow/caffe")
-  parser.add_argument(
-    "--embed_model_data",
-    type=str2bool,
-    default=True,
-    help="input shape.")
-  return parser.parse_known_args()
+    """Parses command line arguments."""
+    parser = argparse.ArgumentParser()
+    parser.register("type", "bool", lambda v: v.lower() == "true")
+    parser.add_argument(
+        "--model_file",
+        type=str,
+        default="",
+        help="TensorFlow \'GraphDef\' file to load, "
+        "Caffe prototxt file to load.")
+    parser.add_argument(
+        "--weight_file", type=str, default="", help="Caffe data file to load.")
+    parser.add_argument(
+        "--model_checksum",
+        type=str,
+        default="",
+        help="Model file sha256 checksum")
+    parser.add_argument(
+        "--weight_checksum",
+        type=str,
+        default="",
+        help="Weight file sha256 checksum")
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="",
+        help="File to save the output graph to.")
+    parser.add_argument(
+        "--runtime", type=str, default="cpu", help="Runtime: cpu/gpu/dsp")
+    parser.add_argument(
+        "--input_node",
+        type=str,
+        default="input_node",
+        help="e.g., input_node")
+    parser.add_argument(
+        "--output_node", type=str, default="softmax", help="e.g., softmax")
+    parser.add_argument(
+        "--data_type",
+        type=str,
+        default='DT_FLOAT',
+        help="e.g., DT_HALF/DT_FLOAT")
+    parser.add_argument(
+        "--output_type", type=str, default="pb", help="output type: source/pb")
+    parser.add_argument(
+        "--template", type=str, default="", help="template path")
+    parser.add_argument(
+        "--obfuscate",
+        type=str2bool,
+        nargs='?',
+        const=False,
+        default=False,
+        help="obfuscate model names")
+    parser.add_argument(
+        "--model_tag",
+        type=str,
+        default="",
+        help="model tag for generated function and namespace")
+    parser.add_argument(
+        "--winograd",
+        type=str2bool,
+        nargs='?',
+        const=False,
+        default=False,
+        help="open winograd convolution or not")
+    parser.add_argument(
+        "--dsp_mode", type=int, default=0, help="dsp run mode, defalut=0")
+    parser.add_argument(
+        "--input_shape", type=str, default="", help="input shape.")
+    parser.add_argument(
+        "--platform", type=str, default="tensorflow", help="tensorflow/caffe")
+    parser.add_argument(
+        "--embed_model_data", type=str2bool, default=True, help="input shape.")
+    return parser.parse_known_args()
 
 
 if __name__ == '__main__':
-  FLAGS, unparsed = parse_args()
-  main(unused_args=[sys.argv[0]] + unparsed)
+    FLAGS, unparsed = parse_args()
+    main(unused_args=[sys.argv[0]] + unparsed)

mace/python/tools/dsp_ops.py

-
 class DspOps(object):
-  def __init__(self):
-    self.dsp_ops = {
-      'INPUT': 'INPUT"',
-      'OUTPUT': 'OUTPUT',
-      'NoOp': 'Nop',
-      'FLATTEN': 'Flatten',
-      'Identity': 'Nop',
-      'Placeholder': 'INPUT',
-      'Const': 'Const',
-      'QuantizedConv2D': 'QuantizedConv2d_8x8to32',
-      'QuantizedMatMul': 'QuantizedMatMul_8x8to32',
-      'QuantizeDownAndShrinkRange': 'QuantizeDownAndShrinkRange_32to8',
-      'QuantizedRelu': 'QuantizedRelu_8',
-      'QuantizedReluX': 'QuantizedReluX_8',
-      'QuantizedMaxPool': 'QuantizedMaxPool_8',
-      'QuantizedAvgPool': 'QuantizedAvgPool_8',
-      'QuantizedConcat': 'QuantizedConcat_8',
-      'QuantizedBiasAdd': 'QuantizedBiasAdd_8p8to32',
-      'QuantizedResizeBilinear' : 'QuantizedResizeBilinear_8',
-      'QuantizedSpaceToBatchND': 'QuantizedSpaceToBatchND_8',
-      'QuantizedBatchToSpaceND': 'QuantizedBatchToSpaceND_8',
-      'QuantizedSoftmax': 'QuantizedSoftmax_8',
-      'QuantizedTanh': 'QuantizedTanh_8',
-      'Min': 'Min_f',
-      'Max': 'Max_f',
-      'QuantizeV2': 'Quantize',
-      'Dequantize': 'Dequantize',
-      'Softmax': 'Softmax_f',
-      'Reshape': 'Reshape',
-      'QuantizedReshape': 'QuantizedReshape',
-      'Sigmoid': 'Sigmoid_f',
-      'Slice': 'Slice_f',
-      'Add': 'Add_f',
-      'Mul': 'Mul_f',
-      'Requantize': 'Requantize_32to8',
-      'RequantizationRange': 'RequantizationRange_32',
-      'Sub': 'Sub_f',
-      'Pack': 'Pack_int32',
-      'StridedSlice': 'StridedSlice_f',
-      'ExpandDims': 'ExpandDims_f',
-      'QuantizedMul': 'QuantizedMul_8x8to32',
-      'QuantizedAdd': 'QuantizedAdd_8p8to32',
-      'Pad': 'Pad_f',
-      'SpaceToBatchND': 'SpaceToBatchND_f',
-      'BatchToSpaceND': 'BatchToSpaceND_f',
-      'ResizeBilinear': 'ResizeBilinear_f',
-      'ConcatV2': 'ConcatV2_f',
-      'Conv2DBackpropInput': 'Deconv_f',
-      'Tanh': 'Tanh_f',
-      'Split': 'Split_f',
-      'Transpose': 'Transpose_f',
-      'Concat': 'Concat_f',
-      'AddN': 'AddN_f',
-    }
-  def has_op(self, tf_op):
-    return tf_op in self.dsp_ops
-
-  def map_nn_op(self, tf_op):
-    if tf_op not in self.dsp_ops:
-      raise Exception('Could not map nn op for: ', tf_op)
-    return self.dsp_ops[tf_op]
+    def __init__(self):
+        self.dsp_ops = {
+            'INPUT': 'INPUT"',
+            'OUTPUT': 'OUTPUT',
+            'NoOp': 'Nop',
+            'FLATTEN': 'Flatten',
+            'Identity': 'Nop',
+            'Placeholder': 'INPUT',
+            'Const': 'Const',
+            'QuantizedConv2D': 'QuantizedConv2d_8x8to32',
+            'QuantizedMatMul': 'QuantizedMatMul_8x8to32',
+            'QuantizeDownAndShrinkRange': 'QuantizeDownAndShrinkRange_32to8',
+            'QuantizedRelu': 'QuantizedRelu_8',
+            'QuantizedReluX': 'QuantizedReluX_8',
+            'QuantizedMaxPool': 'QuantizedMaxPool_8',
+            'QuantizedAvgPool': 'QuantizedAvgPool_8',
+            'QuantizedConcat': 'QuantizedConcat_8',
+            'QuantizedBiasAdd': 'QuantizedBiasAdd_8p8to32',
+            'QuantizedResizeBilinear': 'QuantizedResizeBilinear_8',
+            'QuantizedSpaceToBatchND': 'QuantizedSpaceToBatchND_8',
+            'QuantizedBatchToSpaceND': 'QuantizedBatchToSpaceND_8',
+            'QuantizedSoftmax': 'QuantizedSoftmax_8',
+            'QuantizedTanh': 'QuantizedTanh_8',
+            'Min': 'Min_f',
+            'Max': 'Max_f',
+            'QuantizeV2': 'Quantize',
+            'Dequantize': 'Dequantize',
+            'Softmax': 'Softmax_f',
+            'Reshape': 'Reshape',
+            'QuantizedReshape': 'QuantizedReshape',
+            'Sigmoid': 'Sigmoid_f',
+            'Slice': 'Slice_f',
+            'Add': 'Add_f',
+            'Mul': 'Mul_f',
+            'Requantize': 'Requantize_32to8',
+            'RequantizationRange': 'RequantizationRange_32',
+            'Sub': 'Sub_f',
+            'Pack': 'Pack_int32',
+            'StridedSlice': 'StridedSlice_f',
+            'ExpandDims': 'ExpandDims_f',
+            'QuantizedMul': 'QuantizedMul_8x8to32',
+            'QuantizedAdd': 'QuantizedAdd_8p8to32',
+            'Pad': 'Pad_f',
+            'SpaceToBatchND': 'SpaceToBatchND_f',
+            'BatchToSpaceND': 'BatchToSpaceND_f',
+            'ResizeBilinear': 'ResizeBilinear_f',
+            'ConcatV2': 'ConcatV2_f',
+            'Conv2DBackpropInput': 'Deconv_f',
+            'Tanh': 'Tanh_f',
+            'Split': 'Split_f',
+            'Transpose': 'Transpose_f',
+            'Concat': 'Concat_f',
+            'AddN': 'AddN_f',
+        }
 
+    def has_op(self, tf_op):
+        return tf_op in self.dsp_ops
 
+    def map_nn_op(self, tf_op):
+        if tf_op not in self.dsp_ops:
+            raise Exception('Could not map nn op for: ', tf_op)
+        return self.dsp_ops[tf_op]

mace/python/tools/encrypt_opencl_codegen.py

@@ -4,77 +4,81 @@ import sys
 
 import jinja2
 
-# python encrypt_opencl_codegen.py --cl_kernel_dir=./mace/kernels/opencl/cl/    \
+# python encrypt_opencl_codegen.py --cl_kernel_dir=./mace/kernels/opencl/cl/  \
 #     --output_path=./mace/codegen/opencl_encrypt/opencl_encrypted_program.cc
 
 FLAGS = None
 
 encrypt_lookup_table = "Xiaomi-AI-Platform-Mace"
 
+
 def encrypt_code(code_str):
-  encrypted_arr = []
-  for i in range(len(code_str)):
-    encrypted_char = hex(ord(code_str[i]) ^ ord(encrypt_lookup_table[i % len(encrypt_lookup_table)]))
-    encrypted_arr.append(encrypted_char)
-  return encrypted_arr
+    encrypted_arr = []
+    for i in range(len(code_str)):
+        encrypted_char = hex(
+            ord(code_str[i]) ^ ord(
+                encrypt_lookup_table[i % len(encrypt_lookup_table)]))
+        encrypted_arr.append(encrypted_char)
+    return encrypted_arr
 
 
 def main(unused_args):
-  if not os.path.exists(FLAGS.cl_kernel_dir):
-    print("Input cl_kernel_dir " + FLAGS.cl_kernel_dir + " doesn't exist!")
-
-  header_code = ""
-  for file_name in os.listdir(FLAGS.cl_kernel_dir):
-    file_path = os.path.join(FLAGS.cl_kernel_dir, file_name)
-    if file_path[-2:] == ".h":
-      f = open(file_path, "r")
-      header_code += f.read()
-
-  encrypted_code_maps = {}
-  for file_name in os.listdir(FLAGS.cl_kernel_dir):
-    file_path = os.path.join(FLAGS.cl_kernel_dir, file_name)
-    if file_path[-3:] == ".cl":
-      f = open(file_path, "r")
-      code_str = ""
-      for line in f.readlines():
-        if "#include <common.h>" in line:
-          code_str += header_code
-        else:
-          code_str += line
-      encrypted_code_arr = encrypt_code(code_str)
-      encrypted_code_maps[file_name[:-3]] = encrypted_code_arr
-
-  env = jinja2.Environment(loader=jinja2.FileSystemLoader(sys.path[0]))
-  cpp_cl_encrypted_kernel = env.get_template('str2vec_maps.cc.jinja2').render(
-      maps=encrypted_code_maps,
-      data_type='unsigned char',
-      variable_name='kEncryptedProgramMap')
-
-  if os.path.isfile(FLAGS.output_path):
-    os.remove(FLAGS.output_path)
-  w_file = open(FLAGS.output_path, "w")
-  w_file.write(cpp_cl_encrypted_kernel)
-  w_file.close()
-
-  print("Generate encrypted opencl source done!")
+    if not os.path.exists(FLAGS.cl_kernel_dir):
+        print("Input cl_kernel_dir " + FLAGS.cl_kernel_dir + " doesn't exist!")
+
+    header_code = ""
+    for file_name in os.listdir(FLAGS.cl_kernel_dir):
+        file_path = os.path.join(FLAGS.cl_kernel_dir, file_name)
+        if file_path[-2:] == ".h":
+            f = open(file_path, "r")
+            header_code += f.read()
+
+    encrypted_code_maps = {}
+    for file_name in os.listdir(FLAGS.cl_kernel_dir):
+        file_path = os.path.join(FLAGS.cl_kernel_dir, file_name)
+        if file_path[-3:] == ".cl":
+            f = open(file_path, "r")
+            code_str = ""
+            for line in f.readlines():
+                if "#include <common.h>" in line:
+                    code_str += header_code
+                else:
+                    code_str += line
+            encrypted_code_arr = encrypt_code(code_str)
+            encrypted_code_maps[file_name[:-3]] = encrypted_code_arr
+
+    env = jinja2.Environment(loader=jinja2.FileSystemLoader(sys.path[0]))
+    cpp_cl_encrypted_kernel = env.get_template(
+        'str2vec_maps.cc.jinja2').render(
+            maps=encrypted_code_maps,
+            data_type='unsigned char',
+            variable_name='kEncryptedProgramMap')
+
+    if os.path.isfile(FLAGS.output_path):
+        os.remove(FLAGS.output_path)
+    w_file = open(FLAGS.output_path, "w")
+    w_file.write(cpp_cl_encrypted_kernel)
+    w_file.close()
+
+    print("Generate encrypted opencl source done!")
 
 
 def parse_args():
-  """Parses command line arguments."""
-  parser = argparse.ArgumentParser()
-  parser.add_argument(
-      "--cl_kernel_dir",
-      type=str,
-      default="./mace/kernels/opencl/cl/",
-      help="The cl kernels directory.")
-  parser.add_argument(
-      "--output_path",
-      type=str,
-      default="./mace/examples/codegen/opencl/opencl_encrypted_program.cc",
-      help="The path of encrypted opencl kernels.")
-  return parser.parse_known_args()
+    """Parses command line arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cl_kernel_dir",
+        type=str,
+        default="./mace/kernels/opencl/cl/",
+        help="The cl kernels directory.")
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="./mace/examples/codegen/opencl/opencl_encrypted_program.cc",
+        help="The path of encrypted opencl kernels.")
+    return parser.parse_known_args()
 
 
 if __name__ == '__main__':
-  FLAGS, unparsed = parse_args()
-  main(unused_args=[sys.argv[0]] + unparsed)
+    FLAGS, unparsed = parse_args()
+    main(unused_args=[sys.argv[0]] + unparsed)

mace/python/tools/graph_util.py

@@ -2,18 +2,21 @@ import tensorflow as tf
 from mace.proto import mace_pb2
 from collections import OrderedDict
 
+
 def sort_tf_node(node, nodes_map, ordered_nodes_map):
     if node.name not in ordered_nodes_map:
         for input_tensor_name in node.input:
             input_node_name = input_tensor_name.split(':')[
                 0] if ':' in input_tensor_name else input_tensor_name
-            if input_node_name not in nodes_map or input_node_name in ordered_nodes_map:
+            if input_node_name not in nodes_map or \
+                    input_node_name in ordered_nodes_map:
                 continue
 
             input_node = nodes_map[input_node_name]
             sort_tf_node(input_node, nodes_map, ordered_nodes_map)
         ordered_nodes_map[node.name] = node
 
+
 def sort_tf_graph(graph_def):
     nodes_map = {}
     ordered_nodes_map = OrderedDict()
@@ -31,13 +34,15 @@ def sort_mace_node(node, nodes_map, ordered_nodes_map):
         for input_tensor_name in node.input:
             input_node_name = input_tensor_name.split(':')[
                 0] if ':' in input_tensor_name else input_tensor_name
-            if input_node_name not in nodes_map or input_node_name in ordered_nodes_map:
+            if input_node_name not in nodes_map or \
+                    input_node_name in ordered_nodes_map:
                 continue
 
             input_node = nodes_map[input_node_name]
             sort_mace_node(input_node, nodes_map, ordered_nodes_map)
         ordered_nodes_map[node.name] = node
 
+
 def sort_mace_graph(graph_def, output_name):
     nodes_map = {}
     ordered_nodes_map = OrderedDict()

mace/python/tools/memory_optimizer.py

@@ -2,120 +2,131 @@ import sys
 import operator
 from mace.proto import mace_pb2
 
+
 class MemoryOptimizer(object):
-  def __init__(self, net_def):
-    self.net_def = net_def
-    self.idle_mem = set()
-    self.op_mem = {}    # op_name->mem_id
-    self.mem_block = {} # mem_id->[x, y]
-    self.total_mem_count = 0
-    self.ref_counter = {}
-
-    consumers = {}
-    for op in net_def.op:
-      if self.is_buffer_image_op(op):
-        continue
-      for ipt in op.input:
-        if ipt not in consumers:
-          consumers[ipt] = []
-        consumers[ipt].append(op)
-    # only ref op's output tensor
-    for op in net_def.op:
-      if self.is_buffer_image_op(op):
-        continue
-      for output in op.output:
-        tensor_name = output
-        if tensor_name in consumers:
-          self.ref_counter[tensor_name] = len(consumers[tensor_name])
+    def __init__(self, net_def):
+        self.net_def = net_def
+        self.idle_mem = set()
+        self.op_mem = {}  # op_name->mem_id
+        self.mem_block = {}  # mem_id->[x, y]
+        self.total_mem_count = 0
+        self.ref_counter = {}
+
+        consumers = {}
+        for op in net_def.op:
+            if self.is_buffer_image_op(op):
+                continue
+            for ipt in op.input:
+                if ipt not in consumers:
+                    consumers[ipt] = []
+                consumers[ipt].append(op)
+        # only ref op's output tensor
+        for op in net_def.op:
+            if self.is_buffer_image_op(op):
+                continue
+            for output in op.output:
+                tensor_name = output
+                if tensor_name in consumers:
+                    self.ref_counter[tensor_name] = len(consumers[tensor_name])
+                else:
+                    self.ref_counter[tensor_name] = 0
+
+    def is_buffer_image_op(self, op):
+        return op.type == 'BufferToImage' or op.type == 'ImageToBuffer'
+
+    def get_mem_size(self, op_type, output_shape):
+        mem_size = [0, 0]
+        if op_type == 'WinogradTransform' or op_type == 'MatMul':
+            mem_size[0] = output_shape[2] * output_shape[3]
+            mem_size[1] = output_shape[0] * int((output_shape[1] + 3) / 4)
         else:
-          self.ref_counter[tensor_name] = 0
-
-  def is_buffer_image_op(self, op):
-    return op.type == 'BufferToImage' or op.type == 'ImageToBuffer'
-
-  def get_mem_size(self, op_type, output_shape):
-    mem_size = [0, 0]
-    if op_type == 'WinogradTransform' or op_type == 'MatMul':
-      mem_size[0] = output_shape[2] * output_shape[3]
-      mem_size[1] = output_shape[0] * int((output_shape[1]+3)/4)
-    else:
-      mem_size[0] = output_shape[2] * int((output_shape[3]+3)/4)
-      mem_size[1] = output_shape[0] * output_shape[1]
-    return mem_size
-
-  def mem_area(self, memory_size):
-    return memory_size[0] * memory_size[1]
-
-  def optimize(self):
-    for op in self.net_def.op:
-      if self.is_buffer_image_op(op):
-        continue
-      if not op.output_shape:
-        print('WARNING: There is no output shape information to do memory optimization.')
-        return
-      if len(op.output_shape) != len(op.output):
-        print('WARNING: the number of output shape is not equal to the number of output.')
-        return
-      for i in range(len(op.output)):
-        op_mem_size = self.get_mem_size(op.type, op.output_shape[i].dims)
-        mem_id = -1
-        if len(self.idle_mem) > 0:
-          best_mem_candidate_id = -1
-          best_mem_candidate_delta_area = sys.maxint
-          best_mem_candidate_shape = []
-          for mid in self.idle_mem:
-            reuse_mem_size = self.mem_block[mid]
-            resize_mem_size = [max(reuse_mem_size[0], op_mem_size[0]), max(reuse_mem_size[1], op_mem_size[1])]
-            delta_mem_area = self.mem_area(resize_mem_size) - self.mem_area(reuse_mem_size)
-            if delta_mem_area < best_mem_candidate_delta_area:
-              best_mem_candidate_id = mid
-              best_mem_candidate_delta_area = delta_mem_area
-              best_mem_candidate_shape = resize_mem_size
-
-          if best_mem_candidate_delta_area <= self.mem_area(op_mem_size):
-            # reuse
-            self.mem_block[best_mem_candidate_id] = best_mem_candidate_shape
-            mem_id = best_mem_candidate_id
-            self.idle_mem.remove(mem_id)
-
-        if mem_id == -1:
-          mem_id = self.total_mem_count
-          self.total_mem_count += 1
-          self.mem_block[mem_id] = op_mem_size
-
-        op.mem_id.extend([mem_id])
-        self.op_mem[op.output[i]] = mem_id
-
-      # de-ref input tensor mem
-      for ipt in op.input:
-        if ipt in self.ref_counter:
-          self.ref_counter[ipt] -= 1
-          if self.ref_counter[ipt] == 0:
-            self.idle_mem.add(self.op_mem[ipt])
-          elif self.ref_counter[ipt] < 0:
-            raise Exception('ref count is less than 0')
-
-    for mem in self.mem_block:
-      arena = self.net_def.mem_arena
-      block = arena.mem_block.add()
-      block.mem_id = mem
-      block.x = self.mem_block[mem][0]
-      block.y = self.mem_block[mem][1]
-
-    print('total op: %d', len(self.net_def.op))
-    origin_mem_size = 0
-    optimized_mem_size = 0
-    for op in self.net_def.op:
-      if self.is_buffer_image_op(op):
-        continue
-      origin_mem_size += reduce(operator.mul, op.output_shape[0].dims, 1)
-    for mem in self.mem_block:
-      print mem, self.mem_block[mem]
-      optimized_mem_size += reduce(operator.mul, self.mem_block[mem], 4)
-
-    print('origin mem: %d, optimized mem: %d', origin_mem_size, optimized_mem_size)
+            mem_size[0] = output_shape[2] * int((output_shape[3] + 3) / 4)
+            mem_size[1] = output_shape[0] * output_shape[1]
+        return mem_size
+
+    def mem_area(self, memory_size):
+        return memory_size[0] * memory_size[1]
+
+    def optimize(self):
+        for op in self.net_def.op:
+            if self.is_buffer_image_op(op):
+                continue
+            if not op.output_shape:
+                print('WARNING: There is no output shape information to '
+                      'do memory optimization.')
+                return
+            if len(op.output_shape) != len(op.output):
+                print('WARNING: the number of output shape is not equal to '
+                      'the number of output.')
+                return
+            for i in range(len(op.output)):
+                op_mem_size = self.get_mem_size(op.type,
+                                                op.output_shape[i].dims)
+                mem_id = -1
+                if len(self.idle_mem) > 0:
+                    best_mem_candidate_id = -1
+                    best_mem_candidate_delta_area = sys.maxint
+                    best_mem_candidate_shape = []
+                    for mid in self.idle_mem:
+                        reuse_mem_size = self.mem_block[mid]
+                        resize_mem_size = [
+                            max(reuse_mem_size[0], op_mem_size[0]),
+                            max(reuse_mem_size[1], op_mem_size[1])
+                        ]
+                        delta_mem_area = self.mem_area(
+                            resize_mem_size) - self.mem_area(reuse_mem_size)
+                        if delta_mem_area < best_mem_candidate_delta_area:
+                            best_mem_candidate_id = mid
+                            best_mem_candidate_delta_area = delta_mem_area
+                            best_mem_candidate_shape = resize_mem_size
+
+                    if best_mem_candidate_delta_area <= self.mem_area(
+                            op_mem_size):
+                        # reuse
+                        self.mem_block[
+                            best_mem_candidate_id] = best_mem_candidate_shape
+                        mem_id = best_mem_candidate_id
+                        self.idle_mem.remove(mem_id)
+
+                if mem_id == -1:
+                    mem_id = self.total_mem_count
+                    self.total_mem_count += 1
+                    self.mem_block[mem_id] = op_mem_size
+
+                op.mem_id.extend([mem_id])
+                self.op_mem[op.output[i]] = mem_id
+
+            # de-ref input tensor mem
+            for ipt in op.input:
+                if ipt in self.ref_counter:
+                    self.ref_counter[ipt] -= 1
+                    if self.ref_counter[ipt] == 0:
+                        self.idle_mem.add(self.op_mem[ipt])
+                    elif self.ref_counter[ipt] < 0:
+                        raise Exception('ref count is less than 0')
+
+        for mem in self.mem_block:
+            arena = self.net_def.mem_arena
+            block = arena.mem_block.add()
+            block.mem_id = mem
+            block.x = self.mem_block[mem][0]
+            block.y = self.mem_block[mem][1]
+
+        print('total op: %d', len(self.net_def.op))
+        origin_mem_size = 0
+        optimized_mem_size = 0
+        for op in self.net_def.op:
+            if self.is_buffer_image_op(op):
+                continue
+            origin_mem_size += reduce(operator.mul, op.output_shape[0].dims, 1)
+        for mem in self.mem_block:
+            print mem, self.mem_block[mem]
+            optimized_mem_size += reduce(operator.mul, self.mem_block[mem], 4)
+
+        print('origin mem: %d, optimized mem: %d', origin_mem_size,
+              optimized_mem_size)
 
 
 def optimize_memory(net_def):
-  mem_optimizer = MemoryOptimizer(net_def)
-  mem_optimizer.optimize()
+    mem_optimizer = MemoryOptimizer(net_def)
+    mem_optimizer.optimize()

mace/python/tools/opencl_codegen.py

@@ -14,86 +14,89 @@ FLAGS = None
 
 
 def generate_cpp_source():
-  maps = {}
-  platform_info = ''
-  binary_dirs = FLAGS.cl_binary_dirs.strip().split(",")
-  for binary_dir in binary_dirs:
-    binary_path = os.path.join(binary_dir, FLAGS.built_kernel_file_name)
-    if not os.path.exists(binary_path):
-      continue
-
-    print 'generate opencl code from', binary_path
-    with open(binary_path, "rb") as f:
-      binary_array = np.fromfile(f, dtype=np.uint8)
-
-    idx = 0
-    size, = struct.unpack("Q", binary_array[idx:idx+8])
-    idx += 8
-    for _ in xrange(size):
-      key_size, = struct.unpack("i", binary_array[idx:idx+4])
-      idx += 4
-      key, = struct.unpack(str(key_size) + "s", binary_array[idx:idx+key_size])
-      idx += key_size
-      value_size, = struct.unpack("i", binary_array[idx:idx+4])
-      idx += 4
-      maps[key] = []
-      value = struct.unpack(str(value_size) + "B",
-                            binary_array[idx:idx+value_size])
-      idx += value_size
-      for ele in value:
-        maps[key].append(hex(ele))
-
-    cl_platform_info_path = os.path.join(binary_dir, FLAGS.platform_info_file_name)
-    with open(cl_platform_info_path, 'r') as f:
-      curr_platform_info = f.read()
-    if platform_info != "":
-      assert(curr_platform_info == platform_info)
-    platform_info = curr_platform_info
-
-  env = jinja2.Environment(loader=jinja2.FileSystemLoader(sys.path[0]))
-  return env.get_template('opencl_compiled_kernel.cc.jinja2').render(
-    maps = maps,
-    data_type = 'unsigned char',
-    variable_name = 'kCompiledProgramMap',
-    platform_info = platform_info,
-  )
+    maps = {}
+    platform_info = ''
+    binary_dirs = FLAGS.cl_binary_dirs.strip().split(",")
+    for binary_dir in binary_dirs:
+        binary_path = os.path.join(binary_dir, FLAGS.built_kernel_file_name)
+        if not os.path.exists(binary_path):
+            continue
+
+        print 'generate opencl code from', binary_path
+        with open(binary_path, "rb") as f:
+            binary_array = np.fromfile(f, dtype=np.uint8)
+
+        idx = 0
+        size, = struct.unpack("Q", binary_array[idx:idx + 8])
+        idx += 8
+        for _ in xrange(size):
+            key_size, = struct.unpack("i", binary_array[idx:idx + 4])
+            idx += 4
+            key, = struct.unpack(
+                str(key_size) + "s", binary_array[idx:idx + key_size])
+            idx += key_size
+            value_size, = struct.unpack("i", binary_array[idx:idx + 4])
+            idx += 4
+            maps[key] = []
+            value = struct.unpack(
+                str(value_size) + "B", binary_array[idx:idx + value_size])
+            idx += value_size
+            for ele in value:
+                maps[key].append(hex(ele))
+
+        cl_platform_info_path = os.path.join(binary_dir,
+                                             FLAGS.platform_info_file_name)
+        with open(cl_platform_info_path, 'r') as f:
+            curr_platform_info = f.read()
+        if platform_info != "":
+            assert (curr_platform_info == platform_info)
+        platform_info = curr_platform_info
+
+    env = jinja2.Environment(loader=jinja2.FileSystemLoader(sys.path[0]))
+    return env.get_template('opencl_compiled_kernel.cc.jinja2').render(
+        maps=maps,
+        data_type='unsigned char',
+        variable_name='kCompiledProgramMap',
+        platform_info=platform_info,
+    )
+
 
 def main(unused_args):
 
-  cpp_cl_binary_source = generate_cpp_source()
-  if os.path.isfile(FLAGS.output_path):
-    os.remove(FLAGS.output_path)
-  w_file = open(FLAGS.output_path, "w")
-  w_file.write(cpp_cl_binary_source)
-  w_file.close()
+    cpp_cl_binary_source = generate_cpp_source()
+    if os.path.isfile(FLAGS.output_path):
+        os.remove(FLAGS.output_path)
+    w_file = open(FLAGS.output_path, "w")
+    w_file.write(cpp_cl_binary_source)
+    w_file.close()
 
 
 def parse_args():
-  """Parses command line arguments."""
-  parser = argparse.ArgumentParser()
-  parser.add_argument(
-      "--cl_binary_dirs",
-      type=str,
-      default="",
-      help="The cl binaries directories.")
-  parser.add_argument(
-      "--built_kernel_file_name",
-      type=str,
-      default="",
-      help="The cl binaries directories.")
-  parser.add_argument(
-      "--platform_info_file_name",
-      type=str,
-      default="",
-      help="The cl binaries directories.")
-  parser.add_argument(
-      "--output_path",
-      type=str,
-      default="./mace/examples/codegen/opencl/opencl_compiled_program.cc",
-      help="The path of generated C++ header file which contains cl binaries.")
-  return parser.parse_known_args()
+    """Parses command line arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cl_binary_dirs",
+        type=str,
+        default="",
+        help="The cl binaries directories.")
+    parser.add_argument(
+        "--built_kernel_file_name",
+        type=str,
+        default="",
+        help="The cl binaries directories.")
+    parser.add_argument(
+        "--platform_info_file_name",
+        type=str,
+        default="",
+        help="The cl binaries directories.")
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="./mace/examples/codegen/opencl/opencl_compiled_program.cc",
+        help="The path of generated C++ header file for cl binaries.")
+    return parser.parse_known_args()
 
 
 if __name__ == '__main__':
-  FLAGS, unparsed = parse_args()
-  main(unused_args=[sys.argv[0]] + unparsed)
+    FLAGS, unparsed = parse_args()
+    main(unused_args=[sys.argv[0]] + unparsed)

mace/python/tools/source_converter_lib.py

@@ -6,182 +6,196 @@ import hashlib
 from mace.proto import mace_pb2
 from jinja2 import Environment, FileSystemLoader
 
-
 GENERATED_NAME = set()
 
+
 def generate_obfuscated_name(namespace, name):
-  md5 = hashlib.md5()
-  md5.update(namespace)
-  md5.update(name)
-  md5_digest = md5.hexdigest()
-
-  name = md5_digest[:8]
-  while name in GENERATED_NAME:
-    name = md5_digest
-    assert name not in GENERATED_NAME
-  GENERATED_NAME.add(name)
-  return name
+    md5 = hashlib.md5()
+    md5.update(namespace)
+    md5.update(name)
+    md5_digest = md5.hexdigest()
+
+    name = md5_digest[:8]
+    while name in GENERATED_NAME:
+        name = md5_digest
+        assert name not in GENERATED_NAME
+    GENERATED_NAME.add(name)
+    return name
+
 
 def generate_tensor_map(tensors):
-  tensor_map = {}
-  for t in tensors:
-    if not tensor_map.has_key(t.name):
-        tensor_map[t.name] = generate_obfuscated_name("tensor", t.name)
-  return tensor_map
+    tensor_map = {}
+    for t in tensors:
+        if t.name not in tensor_map:
+            tensor_map[t.name] = generate_obfuscated_name("tensor", t.name)
+    return tensor_map
+
 
 def generate_in_out_map(ops, tensor_map):
-  in_out_map = {}
-  for op in ops:
-    op.name = generate_obfuscated_name("op", op.name)
-    for input_name in op.input:
-        if not in_out_map.has_key(input_name):
-          if tensor_map.has_key(input_name):
-            in_out_map[input_name] = tensor_map[input_name]
-          else:
-            in_out_map[input_name] = generate_obfuscated_name("in", input_name)
-    for output_name in op.output:
-      if not in_out_map.has_key(output_name):
-        if tensor_map.has_key(output_name):
-          in_out_map[output_name] = tensor_map[output_name]
-        else:
-          in_out_map[output_name] = generate_obfuscated_name("out", output_name)
-  return in_out_map
+    in_out_map = {}
+    for op in ops:
+        op.name = generate_obfuscated_name("op", op.name)
+        for input_name in op.input:
+            if input_name not in in_out_map:
+                if input_name in tensor_map:
+                    in_out_map[input_name] = tensor_map[input_name]
+                else:
+                    in_out_map[input_name] = generate_obfuscated_name(
+                        "in", input_name)
+        for output_name in op.output:
+            if output_name not in in_out_map:
+                if output_name in tensor_map:
+                    in_out_map[output_name] = tensor_map[output_name]
+                else:
+                    in_out_map[output_name] = generate_obfuscated_name(
+                        "out", output_name)
+    return in_out_map
+
 
 def obfuscate_name(net_def):
-  input_node = "mace_input_node"
-  output_node = "mace_output_node"
-  tensor_map = generate_tensor_map(net_def.tensors)
-  in_out_map = generate_in_out_map(net_def.op, tensor_map)
-  for t in net_def.tensors:
-    if input_node not in t.name and output_node not in t.name:
-      t.name = tensor_map[t.name]
-  for op in net_def.op:
-    for i in range(len(op.input)):
-      if input_node not in op.input[i]:
-        op.input[i] = in_out_map[op.input[i]]
-    for i in range(len(op.output)):
-      if output_node not in op.output[i]:
-        op.output[i] = in_out_map[op.output[i]]
+    input_node = "mace_input_node"
+    output_node = "mace_output_node"
+    tensor_map = generate_tensor_map(net_def.tensors)
+    in_out_map = generate_in_out_map(net_def.op, tensor_map)
+    for t in net_def.tensors:
+        if input_node not in t.name and output_node not in t.name:
+            t.name = tensor_map[t.name]
+    for op in net_def.op:
+        for i in range(len(op.input)):
+            if input_node not in op.input[i]:
+                op.input[i] = in_out_map[op.input[i]]
+        for i in range(len(op.output)):
+            if output_node not in op.output[i]:
+                op.output[i] = in_out_map[op.output[i]]
+
 
 def rename_tensor(net_def):
-  tensor_map = {}
-  for t in net_def.tensors:
-    if not tensor_map.has_key(t.name):
-      tensor_map[t.name] = "_" + t.name[:-2].replace("/", "_")
-      t.name = tensor_map[t.name]
-  for op in net_def.op:
-    for i in range(len(op.input)):
-      if tensor_map.has_key(op.input[i]):
-        op.input[i] = tensor_map[op.input[i]]
-    for i in range(len(op.output)):
-      if tensor_map.has_key(op.output[i]):
-        op.output[i] = tensor_map[op.output[i]]
+    tensor_map = {}
+    for t in net_def.tensors:
+        if t.name not in tensor_map:
+            tensor_map[t.name] = "_" + t.name[:-2].replace("/", "_")
+            t.name = tensor_map[t.name]
+    for op in net_def.op:
+        for i in range(len(op.input)):
+            if op.input[i] in tensor_map:
+                op.input[i] = tensor_map[op.input[i]]
+        for i in range(len(op.output)):
+            if op.output[i] in tensor_map:
+                op.output[i] = tensor_map[op.output[i]]
+
 
 class TensorInfo:
-  def __init__(self, id, t, runtime):
-    self.id = id
-    self.data_type = mace_pb2.DataType.Name(t.data_type)
-    if t.data_type == mace_pb2.DT_FLOAT:
-      if runtime == 'gpu':
-        self.data_type = mace_pb2.DT_HALF
-        self.data = bytearray(np.array(t.float_data).astype(np.float16).tobytes())
-      else:
-        self.data_type = mace_pb2.DT_FLOAT
-        self.data = bytearray(np.array(t.float_data).astype(np.float32).tobytes())
-    elif t.data_type == mace_pb2.DT_INT32:
-      self.data = bytearray(np.array(t.int32_data).astype(np.int32).tobytes())
-    elif t.data_type == mace_pb2.DT_UINT8:
-      self.data = bytearray(np.array(t.int32_data).astype(np.uint8).tolist())
+    def __init__(self, id, t, runtime):
+        self.id = id
+        self.data_type = mace_pb2.DataType.Name(t.data_type)
+        if t.data_type == mace_pb2.DT_FLOAT:
+            if runtime == 'gpu':
+                self.data_type = mace_pb2.DT_HALF
+                self.data = bytearray(
+                    np.array(t.float_data).astype(np.float16).tobytes())
+            else:
+                self.data_type = mace_pb2.DT_FLOAT
+                self.data = bytearray(
+                    np.array(t.float_data).astype(np.float32).tobytes())
+        elif t.data_type == mace_pb2.DT_INT32:
+            self.data = bytearray(
+                np.array(t.int32_data).astype(np.int32).tobytes())
+        elif t.data_type == mace_pb2.DT_UINT8:
+            self.data = bytearray(
+                np.array(t.int32_data).astype(np.uint8).tolist())
+
 
 def stringfy(value):
-  return ', '.join('"{0}"'.format(w) for w in value)
-
-def convert_to_source(net_def, mode_pb_checksum, template_dir, obfuscate, model_tag, output, runtime, embed_model_data):
-  if obfuscate:
-    obfuscate_name(net_def)
-  else:
-    rename_tensor(net_def)
-
-  # Capture our current directory
-  print template_dir
-
-  # Create the jinja2 environment.
-  j2_env = Environment(loader=FileSystemLoader(template_dir), trim_blocks=True)
-  j2_env.filters['stringfy'] = stringfy
-  output_dir = os.path.dirname(output) + '/'
-  # generate tensor source files
-  template_name = 'tensor_source.jinja2'
-  model_data = []
-  offset = 0
-  counter = 0
-  for t in net_def.tensors:
-    tensor_info = TensorInfo(counter, t, runtime)
-    # align
-    if tensor_info.data_type != 'DT_UINT8' and offset % 4 != 0:
-      padding = 4 - offset % 4
-      model_data.extend(bytearray([0] * padding))
-      offset += padding
+    return ', '.join('"{0}"'.format(w) for w in value)
+
+
+def convert_to_source(net_def, mode_pb_checksum, template_dir, obfuscate,
+                      model_tag, output, runtime, embed_model_data):
+    if obfuscate:
+        obfuscate_name(net_def)
+    else:
+        rename_tensor(net_def)
+
+    # Capture our current directory
+    print template_dir
+
+    # Create the jinja2 environment.
+    j2_env = Environment(
+        loader=FileSystemLoader(template_dir), trim_blocks=True)
+    j2_env.filters['stringfy'] = stringfy
+    output_dir = os.path.dirname(output) + '/'
+    # generate tensor source files
+    template_name = 'tensor_source.jinja2'
+    model_data = []
+    offset = 0
+    counter = 0
+    for t in net_def.tensors:
+        tensor_info = TensorInfo(counter, t, runtime)
+        # align
+        if tensor_info.data_type != 'DT_UINT8' and offset % 4 != 0:
+            padding = 4 - offset % 4
+            model_data.extend(bytearray([0] * padding))
+            offset += padding
+        source = j2_env.get_template(template_name).render(
+            tensor_info=tensor_info,
+            tensor=t,
+            tag=model_tag,
+            runtime=runtime,
+            offset=offset,
+        )
+        model_data.extend(tensor_info.data)
+        offset += len(tensor_info.data)
+        with open(output_dir + 'tensor' + str(counter) + '.cc', "wb") as f:
+            f.write(source)
+        counter += 1
+
+    # generate tensor data
+    template_name = 'tensor_data.jinja2'
     source = j2_env.get_template(template_name).render(
-      tensor_info = tensor_info,
-      tensor = t,
-      tag = model_tag,
-      runtime = runtime,
-      offset = offset,
-    )
-    model_data.extend(tensor_info.data)
-    offset += len(tensor_info.data)
-    with open(output_dir + 'tensor' + str(counter) + '.cc', "wb") as f:
-      f.write(source)
-    counter += 1
-
-  # generate tensor data
-  template_name = 'tensor_data.jinja2'
-  source = j2_env.get_template(template_name).render(
-    tag = model_tag,
-    embed_model_data = embed_model_data,
-    model_data_size = offset,
-    model_data = model_data
-  )
-  with open(output_dir + 'tensor_data' + '.cc', "wb") as f:
-    f.write(source)
-  if not embed_model_data:
-    f = open(output_dir + model_tag + '.data', "wb")
-    f.write(bytearray(model_data))
-    f.close()
-
-  # generate op source files
-  template_name = 'operator.jinja2'
-  counter = 0
-  op_size = len(net_def.op)
-  for start in range(0, op_size, 10):
+        tag=model_tag,
+        embed_model_data=embed_model_data,
+        model_data_size=offset,
+        model_data=model_data)
+    with open(output_dir + 'tensor_data' + '.cc', "wb") as f:
+        f.write(source)
+    if not embed_model_data:
+        f = open(output_dir + model_tag + '.data', "wb")
+        f.write(bytearray(model_data))
+        f.close()
+
+    # generate op source files
+    template_name = 'operator.jinja2'
+    counter = 0
+    op_size = len(net_def.op)
+    for start in range(0, op_size, 10):
+        source = j2_env.get_template(template_name).render(
+            start=start,
+            end=min(start + 10, op_size),
+            net=net_def,
+            tag=model_tag,
+            runtime=runtime,
+        )
+        with open(output_dir + 'op' + str(counter) + '.cc', "wb") as f:
+            f.write(source)
+        counter += 1
+
+    # generate model source files
+    template_name = 'model.jinja2'
+    tensors = [
+        TensorInfo(i, net_def.tensors[i], runtime)
+        for i in range(len(net_def.tensors))
+    ]
     source = j2_env.get_template(template_name).render(
-      start = start,
-      end = min(start+10, op_size),
-      net = net_def,
-      tag = model_tag,
-      runtime = runtime,
-    )
-    with open(output_dir + 'op' + str(counter) + '.cc', "wb") as f:
-      f.write(source)
-    counter += 1
-
-  # generate model source files
-  template_name = 'model.jinja2'
-  tensors = [TensorInfo(i, net_def.tensors[i], runtime) for i in range(len(net_def.tensors))]
-  source = j2_env.get_template(template_name).render(
-    tensors = tensors,
-    net = net_def,
-    tag = model_tag,
-    runtime = runtime,
-    model_pb_checksum = mode_pb_checksum
-  )
-  with open(output, "wb") as f:
-    f.write(source)
-
-  # generate model header file
-  template_name = 'model_header.jinja2'
-  source = j2_env.get_template(template_name).render(
-    tag = model_tag,
-  )
-  with open(output_dir + model_tag + '.h', "wb") as f:
-    f.write(source)
+        tensors=tensors,
+        net=net_def,
+        tag=model_tag,
+        runtime=runtime,
+        model_pb_checksum=mode_pb_checksum)
+    with open(output, "wb") as f:
+        f.write(source)
+
+    # generate model header file
+    template_name = 'model_header.jinja2'
+    source = j2_env.get_template(template_name).render(tag=model_tag, )
+    with open(output_dir + model_tag + '.h', "wb") as f:
+        f.write(source)

mace/python/tools/tf_ops_stats.py

@@ -10,148 +10,174 @@ from tensorflow import gfile
 
 FLAGS = None
 
+
 def hist_inc(hist, key):
-  if key in hist:
-    hist[key] += 1
-  else:
-    hist[key] = 1
+    if key in hist:
+        hist[key] += 1
+    else:
+        hist[key] = 1
+
 
 def to_int_list(long_list):
-  int_list = []
-  for value in long_list:
-    int_list.append(int(value))
-  return int_list
+    int_list = []
+    for value in long_list:
+        int_list.append(int(value))
+    return int_list
+
 
 def main(unused_args):
-  if not FLAGS.input or not gfile.Exists(FLAGS.input):
-    print('Input graph file ' + FLAGS.input + ' does not exist!')
-    return -1
-
-  input_graph_def = tf.GraphDef()
-  with gfile.Open(FLAGS.input, 'rb') as f:
-    data = f.read()
-    input_graph_def.ParseFromString(data)
-
-  with tf.Session() as session:
-    with session.graph.as_default() as graph:
-      tf.import_graph_def(input_graph_def, name='')
-  
-    stats = {}
-    ops = graph.get_operations()
-    # extract kernel size for conv_2d
-    tensor_shapes = {}
-    tensor_values = {}
-    print("=========================consts============================")
-    for op in ops:
-      if op.type == 'Const':
-        for output in op.outputs:
-          tensor_name = output.name
-          tensor = output.eval()
-          tensor_shape = list(tensor.shape)
-          tensor_shapes[tensor_name] = tensor_shape
-          print("Const %s: %s, %d" % (tensor_name, tensor_shape, functools.reduce(operator.mul, tensor_shape, 1)))
-          if len(tensor_shape) == 1 and tensor_shape[0] < 10:
-            tensor_values[tensor_name] = list(tensor)
-
-    print("=========================ops============================")
-    for op in ops:
-      if op.type in ['Conv2D']:
-        padding = op.get_attr('padding')
-        strides = to_int_list(op.get_attr('strides'))
-        data_format = op.get_attr('data_format')
-        ksize = 'Unknown'
-        for input in op.inputs:
-          input_name = input.name
-          if input_name.endswith('weights/read:0'):
-            ksize = input.shape.as_list()
-            break
-          if input_name.endswith('weights:0') and input_name in tensor_shapes:
-            ksize = tensor_shapes[input_name]
-            break
-        print('%s(padding=%s, strides=%s, ksize=%s, format=%s) %s => %s' % (op.type, padding, strides, ksize, data_format, op.inputs[0].shape, op.outputs[0].shape))
-        key = '%s(padding=%s, strides=%s, ksize=%s, format=%s)' % (op.type, padding, strides, ksize, data_format)
-        hist_inc(stats, key)
-      elif op.type in ['FusedResizeAndPadConv2D']:
-        padding = op.get_attr('padding')
-        strides = to_int_list(op.get_attr('strides'))
-        resize_align_corners = op.get_attr('resize_align_corners')
-        ksize = 'Unknown'
-        for input in op.inputs:
-          input_name = input.name
-          if input_name.endswith('weights:0') and input_name in tensor_shapes:
-            ksize = tensor_shapes[input_name]
-            break
-        key = '%s(padding=%s, strides=%s, ksize=%s, resize_align_corners=%s)' % (op.type, padding, strides, ksize, resize_align_corners)
-        hist_inc(stats, key)
-      elif op.type in ['ResizeBilinear']:
-        align_corners = op.get_attr('align_corners')
-        size = 'Unknown'
-        for input in op.inputs:
-          input_name = input.name
-          if input_name.endswith('size:0') and input_name in tensor_values:
-            size = tensor_values[input_name]
-            break
-        key = '%s(size=%s, align_corners=%s)' % (op.type, size, align_corners)
-        print(key)
-        hist_inc(stats, key)
-      elif op.type in ['AvgPool', 'MaxPool']:
-        padding = op.get_attr('padding')
-        strides = to_int_list(op.get_attr('strides'))
-        ksize = to_int_list(op.get_attr('ksize'))
-        data_format = op.get_attr('data_format')
-        key = '%s(padding=%s, strides=%s, ksize=%s)' % (op.type, padding, strides, ksize)
-        hist_inc(stats, key)
-      elif op.type in ['SpaceToBatchND', 'BatchToSpaceND']:
-        block_shape = 'Unknown'
-        for input in op.inputs:
-          input_name = input.name
-          if input_name.endswith('block_shape:0') and input_name in tensor_values:
-            block_shape = tensor_values[input_name]
-            break
-        paddings = 'Unknown'
-        for input in op.inputs:
-          input_name = input.name
-          if input_name.endswith('paddings:0') and input_name in tensor_values:
-            paddings = tensor_values[input_name]
-            break
-        crops = 'Unknown'
-        for input in op.inputs:
-          input_name = input.name
-          if input_name.endswith('crops:0') and input_name in tensor_values:
-            paddings = tensor_values[input_name]
-            break
-        if op.type == 'SpaceToBatchND':
-          key = '%s(block_shape=%s, paddings=%s)' % (op.type, block_shape, paddings)
-        else:
-          key = '%s(block_shape=%s, crops=%s)' % (op.type, block_shape, crops)
-        print(key)
-        hist_inc(stats, key)
-      elif op.type == 'Pad':
-        paddings = 'Unknown'
-        for input in op.inputs:
-          input_name = input.name
-          if input_name.endswith('paddings:0') and input_name in tensor_values:
-            paddings = tensor_values[input_name]
-            break
-        key = '%s(paddings=%s)' % (op.type, paddings)
-        hist_inc(stats, key)
-      else:
-        hist_inc(stats, op.type)
-  
-  print("=========================stats============================")
-  for key, value in sorted(six.iteritems(stats)):
-    print('%s: %d' % (key, value))
+    if not FLAGS.input or not gfile.Exists(FLAGS.input):
+        print('Input graph file ' + FLAGS.input + ' does not exist!')
+        return -1
+
+    input_graph_def = tf.GraphDef()
+    with gfile.Open(FLAGS.input, 'rb') as f:
+        data = f.read()
+        input_graph_def.ParseFromString(data)
+
+    with tf.Session() as session:
+        with session.graph.as_default() as graph:
+            tf.import_graph_def(input_graph_def, name='')
+
+        stats = {}
+        ops = graph.get_operations()
+        # extract kernel size for conv_2d
+        tensor_shapes = {}
+        tensor_values = {}
+        print("=========================consts============================")
+        for op in ops:
+            if op.type == 'Const':
+                for output in op.outputs:
+                    tensor_name = output.name
+                    tensor = output.eval()
+                    tensor_shape = list(tensor.shape)
+                    tensor_shapes[tensor_name] = tensor_shape
+                    print("Const %s: %s, %d" %
+                          (tensor_name, tensor_shape,
+                           functools.reduce(operator.mul, tensor_shape, 1)))
+                    if len(tensor_shape) == 1 and tensor_shape[0] < 10:
+                        tensor_values[tensor_name] = list(tensor)
+
+        print("=========================ops============================")
+        for op in ops:
+            if op.type in ['Conv2D']:
+                padding = op.get_attr('padding')
+                strides = to_int_list(op.get_attr('strides'))
+                data_format = op.get_attr('data_format')
+                ksize = 'Unknown'
+                for input in op.inputs:
+                    input_name = input.name
+                    if input_name.endswith('weights/read:0'):
+                        ksize = input.shape.as_list()
+                        break
+                    if input_name.endswith(
+                            'weights:0') and input_name in tensor_shapes:
+                        ksize = tensor_shapes[input_name]
+                        break
+                print(
+                    '%s(padding=%s, strides=%s, ksize=%s, format=%s) %s => %s'
+                    % (op.type, padding, strides, ksize, data_format,
+                       op.inputs[0].shape, op.outputs[0].shape))
+                key = '%s(padding=%s, strides=%s, ksize=%s, format=%s)' % (
+                    op.type, padding, strides, ksize, data_format)
+                hist_inc(stats, key)
+            elif op.type in ['FusedResizeAndPadConv2D']:
+                padding = op.get_attr('padding')
+                strides = to_int_list(op.get_attr('strides'))
+                resize_align_corners = op.get_attr('resize_align_corners')
+                ksize = 'Unknown'
+                for input in op.inputs:
+                    input_name = input.name
+                    if input_name.endswith(
+                            'weights:0') and input_name in tensor_shapes:
+                        ksize = tensor_shapes[input_name]
+                        break
+                key = '%s(padding=%s, strides=%s, ksize=%s, ' \
+                    'resize_align_corners=%s)' % (op.type, padding, strides,
+                                                  ksize, resize_align_corners)
+                hist_inc(stats, key)
+            elif op.type in ['ResizeBilinear']:
+                align_corners = op.get_attr('align_corners')
+                size = 'Unknown'
+                for input in op.inputs:
+                    input_name = input.name
+                    if input_name.endswith(
+                            'size:0') and input_name in tensor_values:
+                        size = tensor_values[input_name]
+                        break
+                key = '%s(size=%s, align_corners=%s)' % (op.type, size,
+                                                         align_corners)
+                print(key)
+                hist_inc(stats, key)
+            elif op.type in ['AvgPool', 'MaxPool']:
+                padding = op.get_attr('padding')
+                strides = to_int_list(op.get_attr('strides'))
+                ksize = to_int_list(op.get_attr('ksize'))
+                data_format = op.get_attr('data_format')
+                key = '%s(padding=%s, strides=%s, ksize=%s)' % (op.type,
+                                                                padding,
+                                                                strides, ksize)
+                hist_inc(stats, key)
+            elif op.type in ['SpaceToBatchND', 'BatchToSpaceND']:
+                block_shape = 'Unknown'
+                for input in op.inputs:
+                    input_name = input.name
+                    if input_name.endswith(
+                            'block_shape:0') and input_name in tensor_values:
+                        block_shape = tensor_values[input_name]
+                        break
+                paddings = 'Unknown'
+                for input in op.inputs:
+                    input_name = input.name
+                    if input_name.endswith(
+                            'paddings:0') and input_name in tensor_values:
+                        paddings = tensor_values[input_name]
+                        break
+                crops = 'Unknown'
+                for input in op.inputs:
+                    input_name = input.name
+                    if input_name.endswith(
+                            'crops:0') and input_name in tensor_values:
+                        paddings = tensor_values[input_name]
+                        break
+                if op.type == 'SpaceToBatchND':
+                    key = '%s(block_shape=%s, paddings=%s)' % (op.type,
+                                                               block_shape,
+                                                               paddings)
+                else:
+                    key = '%s(block_shape=%s, crops=%s)' % (op.type,
+                                                            block_shape, crops)
+                print(key)
+                hist_inc(stats, key)
+            elif op.type == 'Pad':
+                paddings = 'Unknown'
+                for input in op.inputs:
+                    input_name = input.name
+                    if input_name.endswith(
+                            'paddings:0') and input_name in tensor_values:
+                        paddings = tensor_values[input_name]
+                        break
+                key = '%s(paddings=%s)' % (op.type, paddings)
+                hist_inc(stats, key)
+            else:
+                hist_inc(stats, op.type)
+
+    print("=========================stats============================")
+    for key, value in sorted(six.iteritems(stats)):
+        print('%s: %d' % (key, value))
+
 
 def parse_args():
-  '''Parses command line arguments.'''
-  parser = argparse.ArgumentParser()
-  parser.add_argument(
-    '--input',
-    type=str,
-    default='',
-    help='TensorFlow \'GraphDef\' file to load.')
-  return parser.parse_known_args()
+    '''Parses command line arguments.'''
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--input',
+        type=str,
+        default='',
+        help='TensorFlow \'GraphDef\' file to load.')
+    return parser.parse_known_args()
+
 
 if __name__ == '__main__':
-  FLAGS, unparsed = parse_args()
-  main(unused_args=[sys.argv[0]] + unparsed)
+    FLAGS, unparsed = parse_args()
+    main(unused_args=[sys.argv[0]] + unparsed)

tools/falcon_cli.py

-#-*- coding:utf8 -*-
+import json
+import socket
+import itertools
 
-import json, socket, itertools
 
 class FalconCli(object):
-
     def __init__(self, addr, debug=True, buf_size=1000):
         self.socket_ = socket.create_connection(addr)
         self.stream = self.socket_.makefile()
@@ -16,16 +16,19 @@ class FalconCli(object):
         self.stream.close()
 
     @classmethod
-    def connect(cls, server="transfer.falcon.miliao.srv", port=8433, debug=True, buf_size=1000):
+    def connect(cls,
+                server="transfer.falcon.miliao.srv",
+                port=8433,
+                debug=True,
+                buf_size=1000):
         try:
             return FalconCli((server, port), debug, buf_size)
         except socket.error, exc:
-            print "error: connect to %s:%s error: %s" %(server, port, exc)
+            print "error: connect to %s:%s error: %s" % (server, port, exc)
 
     def call(self, name, *params):
-        request = dict(id=next(self.id_counter),
-                    params=list(params),
-                    method=name)
+        request = dict(
+            id=next(self.id_counter), params=list(params), method=name)
         payload = json.dumps(request).encode()
         if self.debug:
             print "--> req:", payload
@@ -49,7 +52,7 @@ class FalconCli(object):
         resp = []
 
         while True:
-            buf = lines[s:s+self.buf_size]
+            buf = lines[s:s + self.buf_size]
             s = s + self.buf_size
             if len(buf) == 0:
                 break
@@ -57,4 +60,3 @@ class FalconCli(object):
             resp.append(r)
 
         return resp
-