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 physiotwin4d.WorkflowTrainPhysicsNeMo(train_manifests, val_manifests, pca_mean_mesh, output_directory, use_template_surface=False, resume_from=None, training_method=None, log_level=20)

Bases: PhysioTwin4DBase

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

The network is supplied as a training method — pass a configured physiotwin4d.TrainPhysicsNeMoMGN or physiotwin4d.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)

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 physiotwin4d.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)

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

Return type:

None

process()

Train the model and write checkpoints, metadata and logs.

Return type:

dict[str, Any]

Returns:

Dict with output_directory, checkpoint, metadata, training_loss and val_rmse_log.

Training methods

class physiotwin4d.TrainPhysicsNeMoBase(log_level=20)

Bases: PhysioTwin4DBase

Base class for a PhysicsNeMo mesh-stage training method.

Not instantiated directly — use physiotwin4d.TrainPhysicsNeMoMGN or physiotwin4d.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)

Initialize the training method with its default hyper-parameters.

Parameters:

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

set_epochs(epochs)

Set the number of training epochs.

Return type:

None

set_batch_size(batch_size)

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

Return type:

None

set_learning_rate(learning_rate)

Set the Adam learning rate.

Return type:

None

build_model(in_features, out_features)

Construct the (uncompiled) network. Implemented by subclasses.

Return type:

Module

setup_inputs(device, template_mesh, template_coords)

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

Return type:

None

forward(model, node_feats, batch_len)

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

Return type:

Tensor

checkpoint_fields()

Return architecture-specific fields to store in the checkpoint.

Return type:

dict

save_artifacts(output_dir)

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, device, epochs, output_dir, template_mesh, template_coords, resume_from=None)

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

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.

• device (device) – Torch device to train on.

• 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)

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 physiotwin4d.TrainPhysicsNeMoMGN(log_level=20)

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)

Initialize the MeshGraphNet training method.

Parameters:

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

set_processor_size(processor_size)

Set the number of message-passing hops.

Return type:

None

set_hidden_dim(hidden_dim)

Set the processor/encoder/decoder hidden dimension.

Return type:

None

set_num_layers(num_layers)

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

Return type:

None

set_num_processor_checkpoint_segments(num_segments)

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)

Construct the (uncompiled) network. Implemented by subclasses.

Return type:

Module

setup_inputs(device, template_mesh, template_coords)

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

Return type:

None

forward(model, node_feats, batch_len)

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

Return type:

Tensor

checkpoint_fields()

Return architecture-specific fields to store in the checkpoint.

Return type:

dict

save_artifacts(output_dir)

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 physiotwin4d.TrainPhysicsNeMoMLP(log_level=20)

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)

Initialize the fully connected training method.

Parameters:

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

set_layer_size(layer_size)

Set the hidden layer width.

Return type:

None

set_num_layers(num_layers)

Set the number of fully connected layers.

Return type:

None

build_model(in_features, out_features)

Construct the (uncompiled) network. Implemented by subclasses.

Return type:

Module

setup_inputs(device, template_mesh, template_coords)

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

Return type:

None

forward(model, node_feats, batch_len)

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

Return type:

Tensor

checkpoint_fields()

Return architecture-specific fields to store in the checkpoint.

Return type:

dict

save_artifacts(output_dir)

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 physiotwin4d 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.

See Also

• The Per-Subject Manifest

• Predicting With a Mesh-Stage Model

• Train a PhysicsNeMo Surrogate