update the model summary tool to check sub functions

tfjs-converter/package.json

@@ -76,6 +76,7 @@
     "make-version": "sh -c ./scripts/make-version",
     "gen-doc": "ts-node -s ./scripts/gen_doc.ts",
     "gen-json": "ts-node -s ./scripts/gen_json.ts",
+    "model-summary": "ts-node -s ./tools/model_summary.ts",
     "pb2json": "ts-node -s ./tools/pb2json_converter.ts",
     "build-pip-package": "yarn gen-json --test && cd python && ./build-pip-package.sh --test /tmp/tfjs-pips",
     "run-python-tests": "yarn gen-json --test && cd python && ./run-python-tests.sh"

tfjs-converter/tools/model_summary.ts

@@ -18,27 +18,39 @@ import * as fs from 'fs';
 
 function summarize(argv: string[]) {
   if (argv.length < 3) {
-    console.log('Usage: ts-node pb2json.ts model_file');
+    console.log('Usage: yarn model-summary model_file');
     return;
   }
 
   const sourcePath = process.argv[2];
   console.log('reading pb model file: ' + sourcePath);
   const rawdata = fs.readFileSync(sourcePath);
-  const nodes: Array<any> = JSON.parse(rawdata.toString())['modelTopology']['node'];
+  const model = JSON.parse(rawdata.toString());
+  if (model.format !== 'graph-model') {
+    console.log('This tool only supports TFJS Graph models.');
+    return;
+  }
+  // tslint:disable-next-line: no-any
+  const nodes: any[] = model['modelTopology']['node'];
+  const library = model['modelTopology']['library'];
+  if (library != null) {
+    const functions = library['function'];
+    // tslint:disable-next-line: no-any
+    functions.forEach((func: any) => nodes.concat(func['nodeDef']));
+  }
 
   const opCount: {[key: string]: number} = {};
-  for (const opNode of nodes) {
+  nodes.forEach(opNode => {
     let count = 0;
     const op = opNode['op'];
     if (opCount[op]) {
-        count = opCount[op];
-    } 
-    opCount[op] = count + 1; 
-  }
+      count = opCount[op];
+    }
+    opCount[op] = count + 1;
+  });
 
   console.log(opCount);
-  console.log('Total ops = ' + nodes.length);
+  console.log(`Total ops = ${nodes.length}`);
 }
 
 summarize(process.argv);