Training a Mesh-Stage Model

WorkflowTrainPhysicsNeMo owns the data side of training - manifests, normalization statistics, lazy datasets, output directories, checkpoints, metadata and logs. The network and its optimization loop live in the training method it drives.

Workflow

class monai_physio.WorkflowTrainPhysicsNeMo(train_manifests, val_manifests, pca_mean_mesh, output_directory, use_template_surface=False, resume_from=None, training_method=None, log_level=20)[source]

Bases: MONAIPhysioBase

Train a PhysicsNeMo mesh-stage model from per-subject manifests.

The network is supplied as a training method - pass a configured monai_physio.TrainPhysicsNeMoMGN or monai_physio.TrainPhysicsNeMoMLP instance as training_method; a default MeshGraphNet method is used when none is given.

__init__(train_manifests, val_manifests, pca_mean_mesh, output_directory, use_template_surface=False, resume_from=None, training_method=None, log_level=20)[source]

Initialize the training workflow.

Parameters:

train_manifests (list[Path]) – Per-subject manifest files for the training set.
val_manifests (list[Path]) – Per-subject manifest files for the validation set (used for intermittent RMSE reporting during training). May be empty to skip validation.
pca_mean_mesh (Path) – PCA template mesh whose point count matches pca_model.json (typically pca_mean.vtu). Its points define the shared node coordinates and - for the MGN - the mesh-graph topology, so a volumetric template trains on volume points and a surface template on surface points. The sibling pca_model.json (if present) is copied into output_directory for inference.
output_directory (Path) – Directory for checkpoints, metadata and logs.
use_template_surface (bool) – Train on the template’s extracted surface instead of its own points. Set this when the PCA model is volumetric but the manifests reference surface meshes.
resume_from (Optional[Path]) – Optional *_stage_model.pt to resume from; its normalization statistics are inherited so the loaded weights stay valid, and a fresh numbered output directory is used.
training_method (Optional[TrainPhysicsNeMoBase]) – Training method instance carrying the network and its hyper-parameters. Defaults to a new monai_physio.TrainPhysicsNeMoMGN.
log_level (int | str) – Logging level. Default: logging.INFO.

Raises:

ValueError – If train_manifests is empty.
FileNotFoundError – If pca_mean_mesh does not exist.
TypeError – If training_method is neither None nor a TrainPhysicsNeMoBase instance.

set_cache_size(cache_max_samples)[source]

Set the RAM cache budget (decoded phase arrays); 0 = unbounded.

Return type:: None

process()[source]

Train the model and write checkpoints, metadata and logs.

Under a distributed launcher every rank calls this and they train one model together; rank 0 alone writes to output_directory, so on the other ranks checkpoint and metadata come back unset.

Return type:: dict[str, Any]
Returns:: Dict with output_directory, checkpoint, metadata, training_loss and val_rmse_log.

Training methods

class monai_physio.TrainPhysicsNeMoBase(log_level=20)[source]

Bases: MONAIPhysioBase

Base class for a PhysicsNeMo mesh-stage training method.

Not instantiated directly - use monai_physio.TrainPhysicsNeMoMGN or monai_physio.TrainPhysicsNeMoMLP. Subclasses implement the network seams (build_model(), setup_inputs(), forward(), checkpoint_fields(), save_artifacts()) and set the class attributes model_tag, architecture_name and _shuffle_points_within_batch.

model_tag: str = 'base'

architecture_name: str = 'base'

__init__(log_level=20)[source]

Initialize the training method with its default hyper-parameters.

Parameters:: log_level (int | str) – Logging level. Default: logging.INFO.

epochs: int

batch_size: int

learning_rate: float

rmse_log_interval: int

loss_log_interval: int

seed: int

grad_clip_norm: float

set_epochs(epochs)[source]

Set the number of training epochs.

Return type:: None

set_batch_size(batch_size)[source]

Set the mini-batch size, measured in (subject, phase) samples.

Return type:: None

set_learning_rate(learning_rate)[source]

Set the Adam learning rate.

Return type:: None

set_grad_clip_norm(grad_clip_norm)[source]

Set the max gradient norm clipped to before each optimizer step.

A cold-start network can produce a huge, fully finite gradient (an energy-based residual like PhysicsInformedMotion has terms the Jacobian clamp does not cover), large enough to poison Adam’s moment estimates in a single step. Clipping bounds that step; it does not catch a non-finite loss, which the training loop skips outright.

Parameters:: grad_clip_norm (float) – Max L2 norm of the gradient, passed to torch.nn.utils.clip_grad_norm_.
Return type:: None

set_compile_incompatible(reason)[source]

Override whether torch.compile is skipped, and why.

A subclass may set self._compile_incompatible automatically when it knows a specific configuration corrupts under Inductor (see TrainPhysicsNeMoPhysicsInformedMotion.set_mechanics()). Call this afterward to override that – pass None to let torch.compile run again, e.g. to re-test whether a newer torch/CUDA build fixed the bug.

Parameters:: reason (Optional[str]) – Human-readable reason torch.compile is skipped, logged in its place; None re-enables the normal compile attempt.
Return type:: None

build_model(in_features, out_features)[source]

Construct the (uncompiled) network. Implemented by subclasses.

Return type:: Module

setup_inputs(device, template_mesh, template_coords)[source]

Prepare any shared per-forward inputs (MGN graph tensors).

Return type:: None

forward(model, node_feats, batch_len)[source]

Run the network for a flattened (batch_len * n_points, F) batch.

Return type:: Tensor

checkpoint_fields()[source]

Return architecture-specific fields to store in the checkpoint.

Return type:: dict

save_artifacts(output_dir)[source]

Save any architecture-specific artifacts (MGN graph tensors).

Called by train() as soon as setup_inputs() has run, so the artifacts are in place for inference from the first intermittent checkpoint rather than only after the last epoch.

Return type:: None

train(train_dataset, val_dataset, stats, context, epochs, output_dir, template_mesh, template_coords, resume_from=None)[source]

Train the network, returning the model and the loss / RMSE logs.

Every rank runs this. Each steps over its own disjoint slice of the samples and the gradients are averaged across ranks, so the effective batch is batch_size * world_size. Only rank 0 writes to disk.

Parameters:

train_dataset (PhaseSampleDataset) – Lazy training samples built by the workflow.
val_dataset (PhaseSampleDataset) – Lazy validation samples; may be empty.
stats (dict) – Normalization statistics computed by the workflow.
context (DistributedContext) – Rank, device and world size of this process.
epochs (int) – Number of epochs to run. The workflow may clamp epochs (for example in test mode), so the effective count is passed in rather than read from the method.
output_dir (Path) – Directory for the intermittent epoch checkpoints.
template_mesh (DataSet) – Shared PCA template mesh, source of the mesh graph.
template_coords (ndarray) – Shared template node coordinates.
resume_from (Optional[Path]) – Optional checkpoint whose weights are loaded before training starts.

Return type:

tuple[Module, list[float], list[dict]]

Returns:

Tuple of (model, training_losses, val_rmse_log).

build_checkpoint(model, stats)[source]

Assemble a self-describing checkpoint (weights + normalization stats).

Both the periodic epoch checkpoints and the final model share this payload so training can resume from - and inference can load - any saved checkpoint, not just the final one.

Return type:: dict[str, Any]

class monai_physio.TrainPhysicsNeMoMGN(log_level=20)[source]

Bases: TrainPhysicsNeMoBase

Train a PhysicsNeMo MeshGraphNet on mesh stages.

The mesh-graph topology and edge features are extracted once from the shared PCA template mesh (surface or volume) and reused for every (subject, phase) sample; PyTorch Geometric batches join disconnected sub-graphs.

model_tag: str = 'mgn'

architecture_name: str = 'physicsnemo.models.meshgraphnet.MeshGraphNet'

__init__(log_level=20)[source]

Initialize the MeshGraphNet training method.

Parameters:: log_level (int | str) – Logging level. Default: logging.INFO.

processor_size: int

hidden_dim: int

num_layers: int

num_processor_checkpoint_segments: int

set_processor_size(processor_size)[source]

Set the number of message-passing hops.

Return type:: None

set_hidden_dim(hidden_dim)[source]

Set the processor/encoder/decoder hidden dimension.

Return type:: None

set_num_layers(num_layers)[source]

Set the MLP layer count inside each encoder/processor/decoder block.

Return type:: None

set_num_processor_checkpoint_segments(num_segments)[source]

Set the gradient-checkpointing segment count for the processor.

Gradient checkpointing recomputes processor activations during the backward pass instead of storing them, trading compute for GPU memory. A large mesh graph makes that trade worthwhile: a 179k-point, 1.07M-edge template peaks near 43 GiB at batch_size 4 without it.

Parameters:: num_segments (int) – Number of checkpointed segments; 0 (the default) disables checkpointing and stores every activation.
Return type:: None

build_model(in_features, out_features)[source]

Construct the (uncompiled) network. Implemented by subclasses.

Return type:: Module

setup_inputs(device, template_mesh, template_coords)[source]

Prepare any shared per-forward inputs (MGN graph tensors).

Return type:: None

forward(model, node_feats, batch_len)[source]

Run the network for a flattened (batch_len * n_points, F) batch.

Return type:: Tensor

checkpoint_fields()[source]

Return architecture-specific fields to store in the checkpoint.

Return type:: dict

save_artifacts(output_dir)[source]

Save any architecture-specific artifacts (MGN graph tensors).

Called by train() as soon as setup_inputs() has run, so the artifacts are in place for inference from the first intermittent checkpoint rather than only after the last epoch.

Return type:: None

class monai_physio.TrainPhysicsNeMoMLP(log_level=20)[source]

Bases: TrainPhysicsNeMoBase

Train a PhysicsNeMo FullyConnected (MLP) on mesh stages.

Each mesh point is an independent training row; batches group several (subject, phase) samples and shuffle points within the batch to retain gradient mixing while still streaming from disk.

model_tag: str = 'mlp'

architecture_name: str = 'physicsnemo.models.mlp.FullyConnected'

__init__(log_level=20)[source]

Initialize the fully connected training method.

Parameters:: log_level (int | str) – Logging level. Default: logging.INFO.

set_layer_size(layer_size)[source]

Set the hidden layer width.

Return type:: None

set_num_layers(num_layers)[source]

Set the number of fully connected layers.

Return type:: None

build_model(in_features, out_features)[source]

Construct the (uncompiled) network. Implemented by subclasses.

Return type:: Module

setup_inputs(device, template_mesh, template_coords)[source]

Prepare any shared per-forward inputs (MGN graph tensors).

Return type:: None

forward(model, node_feats, batch_len)[source]

Run the network for a flattened (batch_len * n_points, F) batch.

Return type:: Tensor

checkpoint_fields()[source]

Return architecture-specific fields to store in the checkpoint.

Return type:: dict

save_artifacts(output_dir)[source]

Save any architecture-specific artifacts (MGN graph tensors).

Called by train() as soon as setup_inputs() has run, so the artifacts are in place for inference from the first intermittent checkpoint rather than only after the last epoch.

Return type:: None

Example

from monai_physio import TrainPhysicsNeMoMGN, WorkflowTrainPhysicsNeMo

method = TrainPhysicsNeMoMGN()
method.set_epochs(1500)
method.set_batch_size(4)
method.set_processor_size(3)      # message-passing hops
method.set_hidden_dim(128)

workflow = WorkflowTrainPhysicsNeMo(
    train_manifests=train_manifests,
    val_manifests=val_manifests,
    pca_mean_mesh=pca_mean_surface_file,
    output_directory=output_dir,
    training_method=method,
)
result = workflow.process()

checkpoint = result["checkpoint"]

Swap TrainPhysicsNeMoMGN for TrainPhysicsNeMoMLP to train the fully connected network instead; nothing else changes.

Notes

The template mesh sets the node domain. pca_mean_mesh defines the points every subject is expressed on, and - for the MeshGraphNet - the graph topology. Pass a .vtu to train on volume points, a .vtp to train on surface points. When the PCA model is volumetric but your manifests reference surfaces, set use_template_surface=True and the workflow trains on the template’s extracted surface.

Target width is inferred. The stored target array’s column count becomes the network’s output size and is recorded in the checkpoint, so inference rebuilds a matching network without being told.

Resuming writes to a fresh directory. Passing resume_from inherits the prior run’s normalization statistics - so the loaded weights stay valid - and writes to a numbered sibling of output_directory. Read the actual location from result["output_directory"].

Streaming. set_cache_size() bounds how many decoded target arrays stay in RAM, so an arbitrarily large training set streams from disk.

Training a Mesh-Stage Model

Workflow

Training methods

Example

Notes

See Also