tensorflow-helpers/spatial-utils.js

Helper functions to use tensorflow in nodejs for transfer learning, image classification, and more

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.getSpatialNodeNames = getSpatialNodeNames;
exports.getLastSpatialNodeName = getLastSpatialNodeName;
exports.getSpatialNodes = getSpatialNodes;
exports.filterSpatialNodesWithUniqueShapes = filterSpatialNodesWithUniqueShapes;
exports.getSpatialNodesWithUniqueShapes = getSpatialNodesWithUniqueShapes;
function getSpatialNodeNames(model) {
    let names = Object.keys(model.executor.graph.nodes);
    // e.g. 'StatefulPartitionedCall/StatefulPartitionedCall/predict/MobilenetV3/expanded_conv_14/project/BatchNorm/FusedBatchNormV3'
    let spatialNodes = names
        .map(name => {
        let match = name.match(/StatefulPartitionedCall\/StatefulPartitionedCall\/predict\/MobilenetV3\/expanded_conv_([\d]+)\/project\/BatchNorm\/FusedBatchNormV3/);
        // starting from 1
        let layer = +match?.[1];
        return { name, layer };
    })
        .filter(node => node.layer >= 1);
    return spatialNodes;
}
function getLastSpatialNodeName(model) {
    let spatialNodes = getSpatialNodeNames(model);
    let lastNode = spatialNodes.sort((a, b) => b.layer - a.layer)[0];
    if (!lastNode)
        throw new Error('No spatial node found');
    return lastNode.name;
}
function getSpatialNodes(args) {
    let { model, tf } = args;
    let spatialNodes = getSpatialNodeNames(model);
    let shape = model.inputs[0].shape.slice(0);
    shape[0] = 1;
    let input = tf.zeros(shape);
    let output = model.execute(input, spatialNodes.map(node => node.name));
    return output.map((tensor, i) => {
        let { name, layer } = spatialNodes[i];
        return {
            name,
            layer,
            shape: tensor.shape,
        };
    });
}
function filterSpatialNodesWithUniqueShapes(spatialNodes) {
    return spatialNodes
        .slice()
        .sort((a, b) => a.layer - b.layer)
        .filter((current, i, nodes) => {
        let next = nodes[i + 1];
        if (!next)
            return true;
        let current_shape = current.shape.join();
        let next_shape = next.shape.join();
        return current_shape != next_shape;
    });
}
function getSpatialNodesWithUniqueShapes(args) {
    let { model, tf } = args;
    let spatialNodes = getSpatialNodes({ model, tf });
    spatialNodes = filterSpatialNodesWithUniqueShapes(spatialNodes);
    return spatialNodes;
}