tfsummaryvis/src/lib/parse_summary.mjs

60 lines
1.7 KiB
JavaScript

"use strict";
export default function parse_summary(text) {
const lines = text.trim().split(/\r?\n/);
const layers_raw = lines.slice(
lines.findIndex(line => line.startsWith("====")) + 1,
lines.findLastIndex(line => line.startsWith("===="))
);
const layers = [];
let acc = [];
for (const line of layers_raw) {
if(line.trim().length == 0) {
if(acc.length === 0) continue;
console.log(acc.map(layer_line => layer_line.substring(65).trim()).join(""));
// Handle parsed item
const result = {
name_raw: acc.map(layer_line => layer_line.substring(0, 32).trim()).join(""),
output_shape: acc.map(layer_line => layer_line.substring(32, 53).trim()).join(""),
params: parseInt(acc.map(layer_line => layer_line.substring(53, 65).trim()).join("")),
connected_to: JSON.parse(acc.map(layer_line => layer_line.substring(65).trim()).join("")
.replace(/'/g, '"'))
.map(connected_name => connected_name.replace(/(\[0\])+$/, ""))
};
result.type = result.name_raw.match(/ \(([^)]+)\)/)[1];
result.name = result.name_raw.split(/\s+/)[0];
layers.push(result);
acc.length = 0;
}
acc.push(line);
}
const edges = [];
for(const layer of layers) {
for(const connection of layer.connected_to) {
edges.push({
from: connection,
to: layer.name
});
}
}
return {
title: lines[0].match(/^Model: "([^"]+)"/)[1],
params: parseInt(lines.find(line => line.search(/^Total params:/) > -1).split(": ")[1].replace(/,/g, "")),
params_trainable: parseInt(lines.find(line => line.search(/^Trainable params:/) > -1).split(": ")[1].replace(/,/g, "")),
params_nontrainable: parseInt(lines.find(line => line.search(/^Non-trainable params:/) > -1).split(": ")[1].replace(/,/g, "")),
layers,
edges
}
}