tfsummaryvis/src/lib/make_graph.mjs

26 lines
956 B
JavaScript

"use strict";
import nomnoml from 'nomnoml';
export default function make_graph(summary) {
const result = [];
result.push("#.title: visual=none title=center,bold body=center,bold");
result.push(`[<title>Model: ${summary.title}|Total params: ${summary.params}|Trainable params: ${summary.params_trainable}|Non-trainable params: ${summary.params_trainable}|${summary.layers.length} layers, ${summary.edges.length} edges]`)
for(const layer of summary.layers) {
const output_shape = layer.output_shape.replace(/\[/, "\\[")
.replace(/\]/, "\\]");
result.push(`[${layer.name}|${layer.type}|Params: ${layer.params}|Output shape: ${output_shape}]`);
}
for(const edge of summary.edges) {
result.push(`[${edge.from}] --> [${edge.to}]`);
}
const nomnoml_source = result.join("\n");
console.info(`[make_graph] nomnoml source:\n${nomnoml_source}`);
const svg = nomnoml.renderSvg(nomnoml_source);
return { svg, source: nomnoml_source };
}