Skip to content

Commit

Permalink
Improve model summary (tensorflow#3519)
Browse files Browse the repository at this point in the history
FEATURE
* fix functions ops summary

* rewrite the cumsum op with parallel algorithm

* address comments

* sorted the ops
  • Loading branch information
pyu10055 authored Jun 26, 2020
1 parent 51318c1 commit 69e5ed5
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions tfjs-converter/tools/model_summary.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,13 +31,14 @@ function summarize(argv: string[]) {
return;
}
// tslint:disable-next-line: no-any
const nodes: any[] = model['modelTopology']['node'];
let nodes: any[] = model['modelTopology']['node'];
const library = model['modelTopology']['library'];
if (library != null) {
const functions = library['function'];

// tslint:disable-next-line: no-any
if (functions != null) {
functions.forEach((func: any) => nodes.concat(func['nodeDef']));
functions.forEach((func: any) => nodes = nodes.concat(func['nodeDef']));
}
}

Expand All @@ -51,7 +52,8 @@ function summarize(argv: string[]) {
opCount[op] = count + 1;
});

console.log(opCount);
const keys = Object.keys(opCount).sort();
keys.forEach(key => console.log(`${key}: ${opCount[key]}`));
console.log(`Total ops = ${nodes.length}`);
}

Expand Down

0 comments on commit 69e5ed5

Please sign in to comment.