Skip to content

API reference

from neural_flow import (visualize_model, trace_model, draw, animate_inputs, animate_model,
                         FlowConfig, FlowResult, summarize_tensor, visualize_tensor)
from neural_flow import sequences          # pan, crossfade, zoom, occlusion_sweep, slices_to_frames

Functions

visualize_model(model, inputs, output=None, *, config=None, interactive=False, return_result=False, **options)

Traces model on inputs and draws the representation flow.

argument meaning
model torch.nn.Module (switched to eval() during the call, then restored)
inputs tensor, tuple/list of tensors (positional arguments), or dict (keyword arguments)
output path; .png, .svg, .pdf or .html
config optional FlowConfig; keyword options override its fields
interactive also write the HTML explorer next to a static output
return_result return (fig, FlowResult) instead of fig

Returns a matplotlib Figure with fig.flow (FlowResult) and fig.flow_figure (FlowFigure) attached.

trace_model(model, inputs, config=None, **options) -> FlowResult

Runs the metadata pass, stage selection, topology, capture and, if enabled, explanations. Does not draw.

draw(result, strategy=None, frame=None) -> FlowFigure

Renders a FlowResult with its configured style. strategy overrides the channel ranking for this drawing. FlowFigure has .fig, .ax, .stage_artists (stage key → artists), .edge_artists and .layout.

Renders one frame per input with stages, channels, per-channel scales and the PCA colour basis held fixed. inputs is a tensor [T, …] (one frame per leading index) or any iterable of per-frame inputs (tensors or dicts). Writes MP4 (needs imageio-ffmpeg) or GIF, and returns the path written. See movies.md.

animate_model(model, inputs, output="network_flow.gif", *, fps=12, frames_per_stage=6, hold_frames=18, config=None, **options) -> str

A single input; stages light up one after another in data-flow order.

neural_flow.volume3d (3-D models)

function does
sliding_window_movie(model, volume, roi, output="sliding_window.mp4", *, overlap=0.25, max_windows=None, fps=4, title=None, subtitle=None, **options) -> str movie of whole-volume inference: one frame per window, the stages following the window, the fused segmentation assembling in the output card
sliding_window_infer(model, x, roi, overlap=0.25, max_mb=1500, callback=None) -> SlidingWindowState Gaussian-weighted sliding-window fusion of every dense output of x [1, C, X, Y, Z]; state.fused(name) returns labels [1, 1, *S] (multi-class) or values (single channel)
patch_grid(spatial, roi, overlap) window start positions, raster order
flatten_result(result, how="max") what flat_3d=True does, on an existing result

See 3-D models.

neural_flow.nnunet (trained nnU-Net v2 models)

function does
load_nnunet(spec, fold=None, checkpoint="checkpoint_final.pth", device="cpu") -> LoadedModel network from plans.json, weights of one fold, deep supervision off; lm.preprocess, lm.roi (patch size), lm.spacing, lm.sw_overlap (0.5), lm.volume_axes ("zyx"), lm.class_names, lm.info
find_results(spec) the Trainer__Plans__configuration folder for a folder, its Dataset… parent, or a TotalSegmentator alias
NNUNetPreprocessor(plans, configuration) nnU-Net's reading, cropping, normalisation and resampling; pre({"image": path})["image"] is [C, z, y, x]
build_network(plans, configuration, dataset, deep_supervision=False) the untrained network (old and new plans formats)
fetch_totalseg(alias), ct_sample_path() download TotalSegmentator's models / example CT once

zoo.load_model("nnunet:DIR[:FOLD]"), zoo.load_model("totalseg") and zoo.load_input(path, lm, whole=True) wrap these. The adapter context manager neural_flow.adapters.inference_view(model) is what switches deep supervision off. See nnU-Net and TotalSegmentator.

neural_flow.sequences

function returns
pan(image, window=224, steps=48, out_size=None, vertical=False) [T, C, window, window] windows sliding across a wide (or tall) image [C, H, W]
crossfade(a, b, steps=32, ease=True) [T, …] morph between two inputs
zoom(image, start=1.0, end=0.35, steps=40, center=None, out_size=None) [T, C, S, S] zoom into center (fractions of H, W)
occlusion_sweep(image, patch=48, stride=24, value=0.0) [T, C, H, W] sliding occluding patch
slices_to_frames(volume, axis=-1, step=1) [T, C, H, W] 2-D frames through a 3-D volume

summarize_tensor(tensor, role, budget_bytes, cfg, ctx=None) -> TensorSummary and visualize_tensor(summary, cfg, strategy=None, volume_axes="xyz") -> Visual

The tensor → summary → RGBA-raster dispatcher, usable on its own. role ∈ {"input", "activation", "output"}. Visual.image is an [H, W, 4] float array.

FlowConfig options

Every option can be passed as a keyword to visualize_model, trace_model, animate_inputs and animate_model, or collected in a FlowConfig(...). Adapters (CNN, U-Net, transformer, 3-D) may adjust defaults you have not set yourself; for example, 3-D models rank channels by variance.

Stage selection

option default meaning
layer_selection "auto" "auto", "all" (every non-trivial leaf module) or a predicate f(ModuleCall) -> bool
layers None explicit selectors: "name", "name#k" (k-th call of a reused module), "re:<regex>", "type:<ClassRegex>"
exclude None selectors removed from the selection
max_stages 10 hard cap on module stages
target_stages None preferred count (8 technical, 6 story, 7 cinematic)
labels None {module name: display label}

Tensor → picture

option default meaning
max_channels 16 channels retained per stage
channel_strategy "energy" energy (mean x²), variance, spread (p90 − p10), mean_abs, even, pca
normalize "channel" channel or stage (shared) percentile normalization
percentiles (1, 99) display range
diverging False True / "auto": zero-centred colour map for signed activations
cmap theme default activation colour map (magma light, inferno dark/black)
volume_mode "auto" 3-D: volume, ortho, montage, projection
volume_style "voxels" 3-D volume rendering: voxels, cutaway, glow
projection "max" max, mean, meanabs for volume_mode="projection"
volume_axes "xyz" xyz (nibabel/MONAI [X, Y, Z]), zyx (nnU-Net / SimpleITK) or dhw (slice stack [D, H, W])
voxel_spacing None voxel size per spatial axis (tensor order), e.g. (1, 1, 3); volumes are drawn to physical scale
flat_3d False True / "max" / "mean": draw 3-D stages as squashed 2-D projections
sliding_window False trace one window, show the whole-volume output fused from all windows (needs roi_size)
roi_size None sliding-window size, e.g. (96, 96, 96)
roi_center foreground centre voxel the traced window is centred on
sw_overlap 0.25 sliding-window overlap
sw_max_mb 1500 memory budget for fused logits (coarser accumulation grid above it)
aux_outputs False keep deep-supervision outputs (nnU-Net's training heads); by default they are switched off
token_mode "auto" grid, heatmap, pca, l2, mean
cls_tokens, patch_grid inferred number of leading special tokens; (rows, cols) of patch tokens
channels_last inferred force [B, H, W, C] interpretation
capture_attention True record attention weights (MHA / SDPA)

Outputs

option default meaning
class_names None list, or {output name: list}
output_types None {output name: softmax \| sigmoid \| multilabel \| multilabel_probs \| regression \| segmentation \| embedding \| raw}
output_interpreter None f(name, tensor) -> OutputView \| dict \| str
top_k 3 entries listed per classification card
input_names None display names for the inputs
batch_index 0 which sample of a batch to show

Explanations (cinematic)

option default meaning
explain None gradient explanations; None means on for cinematic, off otherwise
explain_max_mb None skip explanations whose backward pass needs more activation memory than this; None: 80 % of the free RAM (or GPU memory)
front_page "pca" front page of each stack: pca (all channels as RGB) or channel (strongest channel)

Capture and memory

option default meaning
capture_device "cpu" where summaries are stored
max_capture_mb 500 total budget for retained summaries
max_spatial_2d, max_spatial_3d 128, 48 display resolution caps for activations
max_input_2d, max_input_3d 256, 128 resolution caps for the input image / volume
topology "auto" auto / runtime (dataflow tracing), fx, sequential
random_state 42 reserved; all reductions are deterministic

Layout and appearance

option default meaning
style "technical" technical, story, cinematic
theme "light" light, dark, black (cinematic defaults to black)
layout "horizontal" horizontal, vertical, wrap
figsize auto e.g. (16, 9); long chains wrap into rows to fill it
dpi 200 raster resolution
title, subtitle auto figure text
font_scale 1.0 scale all text
show_skips True draw skip connections
edge_width "features" arrow width ∝ log(features), or constant
unet_layout "auto" U-shaped layout for encoder/decoder networks
overlay_segmentation True overlay 2-D segmentations on the input

Movie internals (set automatically by animate_inputs)

option meaning
force_channels {stage key: [channel ids]} shown in every frame
force_pca {stage key: (loadings [3, C], mean [C])} fixed PCA colour basis
physical_fov physical field of view (x, y, z), computed from voxel_spacing

animate_inputs also accepts frame_hook(t, result) (edit a frame's result before it is drawn) and thumb_fn(t, result) (film-strip image), which sliding_window_movie uses.

FlowResult

attribute content
config the resolved FlowConfig (after adapter defaults)
stages stages in execution order (Stage: key, label, kind, concept, summary, call, is_head)
graph StageGraph: stages, edges (src, dst, kind ∈ main/skip/merge), outputs, topology_source
views output key → OutputView (kind, headline, subline, items, mask, vector)
explanation Explanation: units (receptive fields, dependency maps), gradcam, contributions
rollout attention-rollout matrix (transformers), if captured
trace raw metadata pass: every module call, shapes, dataflow ops
notes warnings and summarization notes
context extras such as the sliding-window geometry (context["sliding_window"])
summary_table() text table of stages and outputs

Extending

  • Adapters: subclass neural_flow.adapters.Adapter (match, defaults, concepts, and optionally select(trace, cfg) to propose the stages) and register it with neural_flow.adapters.register_adapter(...) to add defaults, stage choice and conceptual stage names for an architecture family. HybridTransformerUNetAdapter (UNETR / Swin UNETR / UNesT) is an example of a select.
  • Custom output semantics: output_interpreter.
  • Custom stage choice: layer_selection=callable receives each ModuleCall (name, type_name, call_index, in_shapes, out_shapes, n_params, depth, …).