in tensorboard/plugins/graph/tf_graph_common/graph.ts [1015:1274]
export function build(
graphDef: tf_graph_proto.GraphDef,
params: BuildParams,
tracker: ProgressTracker
): Promise<SlimGraph> {
/**
* A dictionary that maps each in-embedding node name to the node
* object.
*/
let inEmbedding: {
[nodeName: string]: OpNode;
} = {};
/**
* A dictionary that maps each out-embedding node name to the node
* object.
*/
let outEmbedding: {
[nodeName: string]: OpNode;
} = {};
/**
* A dictionary that maps each node name to an array of the node's
* out-embedding node label objects.
*/
let outEmbeddings: {
[inputName: string]: OpNode[];
} = {};
let isInEmbeddedPred = getEmbedPredicate(params.inEmbeddingTypes);
let isOutEmbeddedPred = getEmbedPredicate(params.outEmbeddingTypes);
let embeddingNodeNames: string[] = [];
let rawNodes = graphDef.node;
/**
* A list of all the non-embedding node names which appear in the processed
* list of raw nodes. Here we pre-allocate enough room for all the rawNodes,
* even though there will some number of embeddings. The excess array length
* is spliced off later.
*
* Experimentation shows that around 30% of the array will go unused, and
* even for very large networks that amounts to less than 10k spaces.
*/
let nodeNames = new Array<string>(rawNodes.length);
return tf_graph_util
.runAsyncTask(
'Normalizing names',
30,
() => {
let opNodes = new Array<OpNode>(rawNodes.length);
let index = 0;
const processRawNode = (rawNode) => {
let opNode = new OpNodeImpl(rawNode);
if (isInEmbeddedPred(opNode)) {
embeddingNodeNames.push(opNode.name);
inEmbedding[opNode.name] = opNode;
return opNode;
}
if (isOutEmbeddedPred(opNode)) {
embeddingNodeNames.push(opNode.name);
outEmbedding[opNode.name] = opNode;
_.each(opNode.inputs, (input) => {
let inputName = input.name;
outEmbeddings[inputName] = outEmbeddings[inputName] || [];
outEmbeddings[inputName].push(opNode);
});
return opNode;
}
// The node is not an embedding, so add it to the names and nodes
// lists.
opNodes[index] = opNode;
nodeNames[index] = opNode.name;
index++;
return opNode;
};
_.each(rawNodes, processRawNode);
const processFunction = (func: tf_graph_proto.FunctionDef) => {
// Give the function itself a node.
const functionNodeName =
FUNCTION_LIBRARY_NODE_PREFIX + func.signature.name;
// Create an op node for the function. Mark it as part of a
// function library.
processRawNode({
name: functionNodeName,
input: [],
device: '',
op: '',
attr: [],
});
// If the function has inputs, make nodes out of them.
if (func.signature.input_arg) {
// Makes an OpNode out of either an input_arg of a library
// function.
let currentInputIndex = 0;
const processInput = (arg) => {
const opNode = processRawNode({
name: functionNodeName + NAMESPACE_DELIM + arg.name,
input: [],
device: '',
op: 'input_arg',
attr: [
{
key: 'T',
value: {
type: arg.type,
},
},
],
});
opNode.functionInputIndex = currentInputIndex;
currentInputIndex++;
};
// Make nodes for input args of the function. Unfortunately, the
// pbtxt configuration language is not rich enough to
// differentiate between an array with 1 item vs 1 object
// property.
if (func.signature.input_arg['name']) {
// There is only 1 input arg.
processInput(func.signature.input_arg);
} else {
// There are several input args.
_.each(func.signature.input_arg, processInput);
}
}
// Make nodes for output args of the function. Track the names of
// output args within the keys of this object. Unlike the
// input_args, the output_args are already defined within the
// node_defs of the library function.
let currentOutputIndex = 0;
const outputArgNames = {};
// If the function has outputs, make nodes out of them.
if (func.signature.output_arg) {
const processOutput = (arg) => {
outputArgNames[functionNodeName + NAMESPACE_DELIM + arg.name] =
currentOutputIndex;
currentOutputIndex++;
};
if (func.signature.output_arg['name']) {
// There is only 1 output arg.
processOutput(func.signature.output_arg);
} else {
// There are several output args.
_.each(func.signature.output_arg, processOutput);
}
}
_.each(func.node_def, (rawNode) => {
// Prefix with the name of the function so that the graph
// correctly computes the hierarchy (and makes metanodes).
rawNode.name = functionNodeName + '/' + rawNode.name;
if (typeof rawNode.input === 'string') {
rawNode.input = [rawNode.input];
}
const opNode = processRawNode(rawNode);
if (_.isNumber(outputArgNames[rawNode.name])) {
// Mark the node as one of the outputs of the function.
opNode.functionOutputIndex = outputArgNames[rawNode.name];
}
_.each(opNode.inputs, (normalizedInput) => {
normalizedInput.name =
functionNodeName + NAMESPACE_DELIM + normalizedInput.name;
});
});
};
if (graphDef.library && graphDef.library.function) {
// This graph contains functions.
_.each(graphDef.library.function, processFunction);
}
opNodes.splice(index);
nodeNames.splice(index);
return opNodes;
},
tracker,
tb_debug.GraphDebugEventId.NORMALIZING_NAMES
)
.then((opNodes) => {
// Create the graph data structure from the graphlib library.
return tf_graph_util.runAsyncTask(
'Building the data structure',
70,
() => {
let normalizedNameDict = mapStrictHierarchy(
nodeNames,
embeddingNodeNames
);
let graph = new SlimGraph();
// Add the nodes to the graph.
_.each(opNodes, (opNode) => {
let normalizedName = normalizedNameDict[opNode.name] || opNode.name;
graph.nodes[normalizedName] = opNode;
// Check if the node has out-embeddings. If yes, add them to the
// node.
if (opNode.name in outEmbeddings) {
opNode.outEmbeddings = outEmbeddings[opNode.name];
// Normalize the names of the out-embeddings.
_.each(opNode.outEmbeddings, (node) => {
node.name = normalizedNameDict[node.name] || node.name;
});
}
// Update the name of the node.
opNode.name = normalizedName;
});
// Visit each node's inputs to add the edges to the graph. If the
// input
// is an in-embedding, then add it to the node's in-embeddings
// instead.
_.each(opNodes, (opNode) => {
_.each(opNode.inputs, (input, i) => {
let inputName = input.name;
if (inputName in inEmbedding) {
let inEmbedNode = inEmbedding[inputName];
opNode.inEmbeddings.push(inEmbedNode);
// Move the inputs of the in-embedding node into incoming
// edges of
// the main node. E.g. the control dependency of a constant
// node
// should be moved to the op node where the constant is
// embedded.
for (let embedInput of inEmbedNode.inputs) {
addEdgeToGraph(
graph,
normalizedNameDict[embedInput.name] || embedInput.name,
opNode,
embedInput,
params,
i
);
}
} else if (inputName in outEmbedding) {
// Move the inputs of the out-embedding node into inputs of
// the main node where the out-embedding points to.
let outEmbedNode = outEmbedding[inputName];
for (let embedInput of outEmbedNode.inputs) {
addEdgeToGraph(
graph,
normalizedNameDict[embedInput.name] || embedInput.name,
opNode,
input,
params,
i
);
}
} else {
addEdgeToGraph(
graph,
normalizedNameDict[inputName] || inputName,
opNode,
input,
params,
i
);
}
});
});
// Normalize the names of in-embeddings.
_.each(inEmbedding, (node, name) => {
node.name = normalizedNameDict[node.name] || node.name;
});
return graph;
},
tracker,
tb_debug.GraphDebugEventId.BUILD_SLIM_GRAPH
);
});
}