cellmap_flow.inference.graph

The graph a torch model computes, for its output activation (portable.activation).

module_graph(model) reads the graph of what the inference server runs, as portable.activation.output_activation takes it: (nodes, output), each node {"op", "inputs", "outputs"} (a clip’s min and max where they are constants), output the name of the value the model returns. The op names are the format’s own; activation.canonical reads them.

• A TorchScript module (a cellmap-models or Hugging Face export’s model.ts, a fly .ts): its inlined_graph, every submodule’s call inlined. Its constants and attribute reads (the weights) are values no node writes; a list built for aten::cat is the cat’s inputs.

• A torch.export program, or a module with an fx graph of its own (an ExportedProgram’s module(), torch.export.unflatten’s): that graph, each submodule called with a graph of its own inlined.

• Any other module (a script’s, a DaCapo run’s, a fly checkpoint’s network in a Sequential with its sigmoid): a torch.fx symbolic trace, torch.nn’s own modules and TorchScript submodules (a fly .ts inside its Sequential) left whole and read as above. A forward that branches on its input’s values cannot be traced: no graph.

An in-place op (aten::sigmoid_, x.sigmoid_()) changes the value it is given, and a traced graph goes on reading that value by its old name (x.sigmoid_(); return x returns the name the sigmoid was given): each later read of it is read as the in-place op’s output instead.

torch is imported when a graph is first read: the runner imports this module.

Attributes

logger

Functions

module_graph(model)

(nodes, output) of model, a torch module or torch.export program; None when it cannot be read.

Module Contents

cellmap_flow.inference.graph.logger

cellmap_flow.inference.graph.module_graph(model)

(nodes, output) of model, a torch module or torch.export program; None when it cannot be read.