Add caffe ops stat script and fix addn bugs.

mace/kernels/addn.h

@@ -11,10 +11,9 @@
 namespace mace {
 namespace kernels {
 
-struct AddNFunctorBase {};
 
 template <DeviceType D, typename T>
-struct AddNFunctor : AddNFunctorBase {
+struct AddNFunctor {
   void operator()(const std::vector<const Tensor *> &input_tensors,
                   Tensor *output_tensor, StatsFuture *future) {
     output_tensor->ResizeLike(input_tensors[0]);
@@ -24,10 +23,6 @@ struct AddNFunctor : AddNFunctorBase {
     memset(output_ptr, 0, size * sizeof(T));
     int n = input_tensors.size();
     for (int i = 0; i < n; ++i) {
-      MACE_CHECK(input_tensors[i]->dim(0) == output_tensor->dim(0));
-      MACE_CHECK(input_tensors[i]->dim(1) == output_tensor->dim(1));
-      MACE_CHECK(input_tensors[i]->dim(2) == output_tensor->dim(2));
-      MACE_CHECK(input_tensors[i]->dim(3) == output_tensor->dim(3));
       Tensor::MappingGuard input_map(input_tensors[i]);
       const T *input_ptr = input_tensors[i]->data<T>();
       for (index_t j = 0; j < size; ++j) {
@@ -44,7 +39,7 @@ void AddNFunctor<DeviceType::NEON, float>::operator()(
     StatsFuture *future);
 
 template <typename T>
-struct AddNFunctor<DeviceType::OPENCL, T> : AddNFunctorBase {
+struct AddNFunctor<DeviceType::OPENCL, T> {
   void operator()(const std::vector<const Tensor *> &input_tensors,
                   Tensor *output_tensor, StatsFuture *future);
 };

mace/kernels/opencl/addn.cc

@@ -17,6 +17,8 @@ static void AddN(const std::vector<const Tensor *> &input_tensors,
   if (input_tensors.size() > 4) {
     MACE_NOT_IMPLEMENTED;
   }
+  output->ResizeLike(input_tensors[0]);
+
   const index_t batch = output->dim(0);
   const index_t height = output->dim(1);
   const index_t width = output->dim(2);
@@ -36,8 +38,8 @@ static void AddN(const std::vector<const Tensor *> &input_tensors,
 
   uint32_t idx = 0;
   for (auto input : input_tensors) {
-  addn_kernel.setArg(idx++,
-                     *(static_cast<const cl::Image2D *>(input->buffer())));
+    addn_kernel.setArg(idx++,
+                       *(static_cast<const cl::Image2D *>(input->buffer())));
   }
   addn_kernel.setArg(idx++, *(static_cast<cl::Image2D *>(output->buffer())));
 

mace/ops/addn.h

@@ -17,11 +17,14 @@ class AddNOp : public Operator<D, T> {
       : Operator<D, T>(operator_def, ws) {}
 
   bool Run(StatsFuture *future) override {
-    Tensor *output_tensor = this->outputs_[0];
+    Tensor *output_tensor = this->Output(0);
     int n = this->inputs_.size();
     vector<const Tensor *> inputs(n, nullptr);
-    for (int i = 0; i < n; ++i) {
-      inputs[i] = this->inputs_[i];
+    inputs[0] = this->Input(0);
+    for (int i = 1; i < n; ++i) {
+      inputs[i] = this->Input(i);
+      MACE_CHECK(inputs[0]->dim_size() == inputs[i]->dim_size());
+      MACE_CHECK(inputs[0]->size() == inputs[i]->size());
     }
 
     functor_(inputs, output_tensor, future);

mace/proto/BUILD

@@ -18,3 +18,12 @@ py_proto_library(
     srcs_version = "PY2AND3",
     deps = ["@com_google_protobuf//:protobuf_python"],
 )
+
+py_proto_library(
+    name = "caffe_py",
+    srcs = ["caffe.proto"],
+    default_runtime = "@com_google_protobuf//:protobuf_python",
+    protoc = "@com_google_protobuf//:protoc",
+    srcs_version = "PY2AND3",
+    deps = ["@com_google_protobuf//:protobuf_python"],
+)

mace/python/tools/BUILD

@@ -43,3 +43,12 @@ py_binary(
         "//mace/proto:mace_py",
     ],
 )
+
+py_binary(
+    name = "caffe_ops_stats",
+    srcs = ["caffe_ops_stats.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//mace/proto:caffe_py",
+    ],
+)

mace/python/tools/caffe_ops_stats.py

+from mace.proto import caffe_pb2
+import google.protobuf.text_format
+import operator
+import functools
+import argparse
+import sys
+import six
+
+FLAGS = None
+
+def main(unused_args):
+  net = caffe_pb2.NetParameter()
+  with open(FLAGS.input) as f:
+    google.protobuf.text_format.Merge(str(f.read()), net)
+
+  ops = {}
+  for layer in net.layer:
+    if layer.type not in ops:
+      ops[layer.type] = 1
+    else:
+      ops[layer.type] += 1
+
+  for key, value in sorted(ops.items(), key=operator.itemgetter(1)):
+    print key, ":", value
+
+def parse_args():
+  '''Parses command line arguments.'''
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+    '--input',
+    type=str,
+    default='',
+    help='Caffe \'GraphDef\' file to load.')
+  return parser.parse_known_args()
+
+if __name__ == '__main__':
+  FLAGS, unparsed = parse_args()
+  main(unused_args=[sys.argv[0]] + unparsed)

mace/python/tools/tf_converter_lib.py

@@ -465,6 +465,8 @@ class TFConverter(object):
       and self.tf_graph[final_op.name][0].type == 'BatchToSpaceND':
       final_op = self.tf_graph[final_op.name][0]
       self.resolved_ops[final_op.name] = 1
+      self.unused_tensor.add(get_input_tensor(final_op, 1).name)
+      self.unused_tensor.add(get_input_tensor(final_op, 2).name)
     else:
       raise Exception('Convert atrous conv error: no BatchToSpaceND op')