How to: Create and Use a ModelNode

A ModelNode is a computational node that wraps a machine-learning model for use within a ModelGraph. It:

• Wraps a backend-specific model (PyTorch, TensorFlow, or scikit-learn) as a BaseModel

• Receives data from a single upstream source (FeatureSet or another ModelNode)

• Optionally holds an Optimizer for standalone training

• Is composed into a ModelGraph, which is then used in an Experiment

Note: Users typically interact with Experiment at the top level. ModelNode is the building block that ModelGraph orchestrates. This guide covers the full ModelNode API for users who need fine-grained control.

This notebook covers:

• The Model Hierarchy

• Built-In Models

• Wrapping Custom PyTorch Models

• Scikit-Learn Models

• Creating a ModelNode

• The Optimizer

• Building and Running a ModelNode

• Chaining Nodes

• Freezing and Unfreezing

• Serialization

• Summary

import numpy as np
import torch

from modularml import Experiment, FeatureSet, ModelNode, Optimizer

# Note that we don't need to explicitly create an Experiment right away
# We do it here so we can disable the warning raise when creating multiple
# nodes with the same name (`registration_policy` is what controls this).
exp = Experiment(label="create_modelnode", registration_policy="overwrite")

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[1], line 2
      1 import numpy as np
----> 2 import torch
      4 from modularml import Experiment, FeatureSet, ModelNode, Optimizer
      6 # Note that we don't need to explicitly create an Experiment right away
      7 # We do it here so we can disable the warning raise when creating multiple
      8 # nodes with the same name (`registration_policy` is what controls this).

ModuleNotFoundError: No module named 'torch'

We’ll use a simple synthetic dataset throughout this notebook: 500 samples of a 10-point voltage signal with a scalar state-of-health (SOH) target.

rng = np.random.default_rng(42)

fs = FeatureSet.from_dict(
    label="SensorData",
    data={
        "voltage": list(rng.standard_normal((500, 10))),
        "soh": list(rng.standard_normal((500, 1))),
    },
    feature_keys="voltage",
    target_keys="soh",
)
print(fs)

Since FeatureSet can contains more columns than wanted for a certain models inputs, we need to specify which columns are intended to by input to the model.

This is done with the .reference() method on FeatureSets. Going forward, our models will be trained on only the voltage feature, and estimate only the soh target.

fs_ref = fs.reference(features="voltage", targets="soh")
print(fs_ref)

---

The Model Hierarchy

Before creating a ModelNode, it helps to understand the model abstraction layers:

BaseModel (abstract)
├── TorchBaseModel          # Base for built-in PyTorch models
│   ├── SequentialMLP       # Built-in MLP
│   └── SequentialCNN       # Built-in 1D CNN
├── TorchModelWrapper       # Wraps any torch.nn.Module
├── TensorflowModelWrapper  # Wraps any tf.keras.Model
└── ScikitModelWrapper      # Wraps any sklearn.BaseEstimator

All models used in a ModelNode must conform to the BaseModel interface. You can either:

1. Use a built-in model (e.g., SequentialMLP)

2. Wrap your own model with TorchModelWrapper, TensorflowModelWrapper, or ScikitModelWrapper

3. Pass a raw model directly — it will be auto-wrapped via wrap_model()

---

Built-In Models

ModularML provides ready-to-use model architectures in modularml.models. There are currently only built-in models using the PyTorch backend modularml.models.torch. More will be added soon.

These built in models inherit from BaseModel and support lazy shape inference; you can provide shapes at construction time or let ModelGraph.build() infer them automatically.

SequentialMLP

A configurable multi-layer perceptron. Inputs are flattened, passed through n_layers fully-connected layers with activation and optional dropout, then reshaped to output_shape.

Parameter	Type	Default	Description
input_shape	tuple[int, ...]	None	Input shape (no batch dim). Inferred at build if None.
output_shape	tuple[int, ...]	None	Output shape (no batch dim). Inferred at build if None.
n_layers	int	2	Number of linear layers.
hidden_dim	int	32	Hidden units per layer.
activation	str	"relu"	Activation function ("relu", "gelu", "tanh", etc.).
dropout	float	0.0	Dropout rate (0 = no dropout).

from modularml.models.torch import SequentialMLP

# Option A: Provide both shapes up front (builds immediately)
mlp_eager = SequentialMLP(
    input_shape=(1, 10),
    output_shape=(1, 1),
    n_layers=3,
    hidden_dim=64,
    activation="gelu",
    dropout=0.1,
)
print(f"Eager MLP built: {mlp_eager.is_built}")
print(f"  input_shape:  {mlp_eager.input_shape}")
print(f"  output_shape: {mlp_eager.output_shape}")

# Option B: Defer shapes (lazy build - ModelGraph.build() will fill them in)
mlp_lazy = SequentialMLP(
    output_shape=(1, 1),
    n_layers=2,
    hidden_dim=32,
)
print(f"Lazy MLP built: {mlp_lazy.is_built}")
print(f"  input_shape:  {mlp_lazy.input_shape}")
print(f"  output_shape: {mlp_lazy.output_shape}")

SequentialCNN

A 1D convolutional network. Stacks Conv1d layers with optional pooling, dropout, and a final linear projection to output_shape.

Parameter	Type	Default	Description
input_shape	tuple[int, ...]	None	Input shape as (num_channels, length).
output_shape	tuple[int, ...]	None	Output shape (no batch dim).
n_layers	int	2	Number of Conv1d layers.
hidden_dim	int	16	Output channels per Conv1d layer.
kernel_size	int	3	Convolution kernel size.
padding	int	1	Convolution padding.
stride	int	1	Convolution stride.
activation	str	"relu"	Activation function.
dropout	float	0.0	Dropout rate.
pooling	int	1	MaxPool1d kernel size (1 = no pooling).
flatten_output	bool	True	Whether to flatten and project to output shape.

from modularml.models.torch import SequentialCNN

cnn = SequentialCNN(
    input_shape=(1, 10),  # 1 channel, 10-length signal
    output_shape=(1, 1),
    n_layers=2,
    hidden_dim=16,
    kernel_size=3,
    pooling=2,
)
print(f"CNN built: {cnn.is_built}")
print(f"  input_shape:  {cnn.input_shape}")
print(f"  output_shape: {cnn.output_shape}")

---

Wrapping Custom PyTorch Models

For models not provided by ModularML, use TorchModelWrapper to wrap any torch.nn.Module.

Wrapping an Instantiated Model

If you already have a constructed torch.nn.Module, pass it directly. The wrapper validates input/output shapes with a dummy forward pass during build().

from modularml.core.models import TorchModelWrapper


# Define a custom PyTorch model
class MyEncoder(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        # Storing constructor args as same-named attributes
        # allows TorchModelWrapper to auto-infer them for serialization.
        self.in_features = in_features
        self.out_features = out_features
        self.fc = torch.nn.Linear(in_features, out_features)

    def forward(self, x):
        return self.fc(x)


# Wrap an already-instantiated model
raw_model = MyEncoder(in_features=10, out_features=4)
wrapped = TorchModelWrapper(model=raw_model)

print(f"Wrapped model built: {wrapped.is_built}")
print(f"  backend: {wrapped.backend}")

# Build to validate shapes
wrapped.build(input_shape=(10,), output_shape=(4,))
print(f"After build: {wrapped.is_built}")
print(f"  input_shape:  {wrapped.input_shape}")
print(f"  output_shape: {wrapped.output_shape}")

Note that custom models cannot be serialized unless they are defined in a separate Python file.

Example use case:

    from my_scipt import MyModel

    model = MyModel(...)

More details on this are provided in the [Serialization]serialization section.

Lazy Construction from a Class

If you want ModelGraph.build() to handle instantiation (injecting the correct input_shape and output_shape), pass a class and kwargs instead.

lazy_wrapped = TorchModelWrapper(
    model_class=MyEncoder,
    model_kwargs={"in_features": 10, "out_features": 4},
)
print(f"Lazy wrapped built: {lazy_wrapped.is_built}")

# Build later (or let ModelGraph do it)
lazy_wrapped.build(input_shape=(10,), output_shape=(4,))
print(f"After build: {lazy_wrapped.is_built}")
print(f"  output_shape: {lazy_wrapped.output_shape}")

Injecting Shapes into Custom Constructors

By default, TorchModelWrapper injects the inferred input_shape and output_shape into your model class constructor during lazy build. If your constructor uses different parameter names, specify them with inject_input_shape_as and inject_output_shape_as.

class CustomModel(torch.nn.Module):
    """A model whose constructor uses non-standard shape parameter names."""

    def __init__(self, in_shape, out_shape):
        super().__init__()
        self.in_shape = in_shape
        self.out_shape = out_shape
        self.fc = torch.nn.Linear(int(np.prod(in_shape)), int(np.prod(out_shape)))

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x.view(x.size(0), *self.out_shape)


# Tell the wrapper to inject shapes using your parameter names
custom_wrapped = TorchModelWrapper(
    model_class=CustomModel,
    model_kwargs={},
    inject_input_shape_as="in_shape",
    inject_output_shape_as="out_shape",
)
custom_wrapped.build(input_shape=(10,), output_shape=(4,))
print(f"Custom wrapped output_shape: {custom_wrapped.output_shape}")

Auto-Wrapping with wrap_model()

When you pass a raw torch.nn.Module directly to ModelNode, it is automatically wrapped via wrap_model(). This is the simplest path but offers less control over serialization and shape injection.

from modularml.core.models import wrap_model

raw_module = MyEncoder(in_features=10, out_features=4)
auto_wrapped = wrap_model(raw_module)

print(f"Type: {type(auto_wrapped).__name__}")
print(f"Backend: {auto_wrapped.backend}")

---

Scikit-Learn Models

The ScikitModelWrapper wraps any sklearn.base.BaseEstimator for use in a ModelNode. It supports both batch-fit models (e.g., RandomForestRegressor) and incremental models (e.g., SGDRegressor) via the training_mode parameter.

Parameter	Type	Default	Description
model	BaseEstimator	(required)	A scikit-learn estimator instance.
training_mode	str	"auto"	"auto", "partial_fit", or "batch_fit".
output_method	str	"auto"	"auto", "predict", "predict_proba", or "decision_function".
partial_fit_kwargs	dict	None	Extra kwargs passed to every partial_fit() call (e.g., {"classes": [0, 1]}).

Batch-Fit Models

Most scikit-learn models are trained on the full dataset at once. These are used with FitPhase in an Experiment rather than TrainPhase.

from sklearn.ensemble import RandomForestRegressor

from modularml.core.models import ScikitModelWrapper

sklearn_model = ScikitModelWrapper(
    model=RandomForestRegressor(n_estimators=50, random_state=42),
)

print(f"Backend: {sklearn_model.backend}")
print(f"Supports partial_fit: {sklearn_model.supports_partial_fit}")
print(f"Training mode: {sklearn_model.resolved_training_mode}")

Incremental Models

Models that support partial_fit() can be used with TrainPhase for mini-batch training, similar to neural network workflows.

from sklearn.linear_model import SGDRegressor

incremental_model = ScikitModelWrapper(
    model=SGDRegressor(random_state=42),
)

print(f"Supports partial_fit: {incremental_model.supports_partial_fit}")
print(f"Training mode: {incremental_model.resolved_training_mode}")

Auto-Wrapping

Like PyTorch models, raw scikit-learn estimators passed to ModelNode are automatically wrapped via wrap_model().

auto_sklearn = wrap_model(RandomForestRegressor(n_estimators=10))
print(f"Type: {type(auto_sklearn).__name__}")
print(f"Backend: {auto_sklearn.backend}")

---

Creating a ModelNode

A ModelNode combines a model with an upstream data source and an optional optimizer.

    ModelNode(
        label: str,
        model: BaseModel | Any,
        upstream_ref: ExperimentNode | ExperimentNodeReference,
        optimizer: Optimizer | None = None,
    )

Parameter	Description
label	Unique name for this node within the graph.
model	A BaseModel instance, or any raw model (auto-wrapped via wrap_model()).
upstream_ref	The data source: a FeatureSetReference, FeatureSet, or another ModelNode.
optimizer	Optional Optimizer for standalone training. When using ModelGraph, the graph optimizer is typically used instead.

With a Built-In Model

node_mlp = ModelNode(
    label="MyMLP",
    model=SequentialMLP(output_shape=(1, 1), n_layers=2, hidden_dim=32),
    upstream_ref=fs_ref,
)
print(node_mlp)
print(f"  is_built: {node_mlp.is_built}")
print(f"  backend:  {node_mlp.backend}")

With a Custom torch.nn.Module (Auto-Wrapped)

# Pass a raw torch.nn.Module - it is automatically wrapped by wrap_model()
node_custom = ModelNode(
    label="MyCustomEncoder",
    model=MyEncoder(in_features=10, out_features=4),
    upstream_ref=fs_ref,
)
print(node_custom)
print(f"  model type: {type(node_custom.model).__name__}")
print(f"  backend:    {node_custom.backend}")

With a Scikit-Learn Model

node_rf = ModelNode(
    label="RandomForest",
    model=RandomForestRegressor(n_estimators=50, random_state=42),
    upstream_ref=fs_ref,
)
print(node_rf)
print(f"  model type: {type(node_rf.model).__name__}")
print(f"  backend:    {node_rf.backend}")

---

The Optimizer

The Optimizer class wraps backend-specific optimizers with a consistent API.

    Optimizer(
        opt: str | type | None = None,
        *,
        opt_kwargs: dict[str, Any] | None = None,
        factory: Callable | None = None,
        backend: Backend | None = None,
    )

There are three ways to specify the optimizer:

# 1. By name string (most common)
opt_by_name = Optimizer("adam", opt_kwargs={"lr": 1e-3}, backend="torch")

# 2. By optimizer class
opt_by_class = Optimizer(
    torch.optim.AdamW,
    opt_kwargs={"lr": 1e-3, "weight_decay": 1e-4},
)

# 3. By factory callable
opt_by_factory = Optimizer(
    factory=lambda params: torch.optim.SGD(params, lr=0.01, momentum=0.9),
    backend="torch",
)

print(f"By name:    {opt_by_name.name}")
print(f"By class:   {opt_by_class.cls}")
print(f"By factory: {opt_by_factory}")

Attaching an Optimizer to a ModelNode

An optimizer on a ModelNode enables standalone train_step() / eval_step() calls.

If creating a ModelGraph with models all from the same backend (e.g., all PyTorch models), it’s easier to just use a graph-wise optimizer (set during ModelGraph init).

node_with_opt = ModelNode(
    label="TrainableMLP",
    model=SequentialMLP(output_shape=(1, 1), n_layers=2, hidden_dim=32),
    upstream_ref=fs_ref,
    optimizer=Optimizer("adam", opt_kwargs={"lr": 1e-3}, backend="torch"),
)
print(f"Has optimizer: {node_with_opt._optimizer is not None}")

---

Building and Running a ModelNode

Normally ModelGraph.build() handles building all nodes. But for debugging or standalone use, you can build and forward through a ModelNode directly.

Building Manually

# build_model() takes explicit shapes
node_mlp.build_model(input_shape=(1, 10), output_shape=(1, 1))

print(f"is_built:     {node_mlp.is_built}")
print(f"input_shape:  {node_mlp.input_shape}")
print(f"output_shape: {node_mlp.output_shape}")

Forward Pass with SampleData

The forward_single() method (also available as __call__) accepts SampleData, RoleData, or Batch. It passes features through the model, preserving targets and tags.

from modularml.core.data.sample_data import SampleData
from modularml.utils.data.data_format import DataFormat

# Create SampleData from the FeatureSet reference
fsv = fs_ref.resolve()
sample_data = SampleData(
    features=fsv.get_features(fmt=DataFormat.TORCH),
    targets=fsv.get_targets(fmt=DataFormat.TORCH),
)
print(f"Input features: {sample_data.features.shape}")

# Forward pass
with torch.no_grad():
    output = node_mlp(sample_data)
    print(f"Output features: {output.features.shape}")
    print(f"Targets passed through: {output.targets.shape}")

Auto-Build on First Forward Pass

If a ModelNode has a FeatureSetReference as its upstream, calling forward_single() on an unbuilt node will attempt to auto-build by inferring shapes from the upstream FeatureSet.

Note that output_shape will be determined in the following sequence:

• If provided, that output shape is used

• If the node has no downstream connections, the target shape of the referenced FeatureSet will be used

• Otherwise, the hidden layer shape (if using a built model) will be used.

# This node is not built yet
auto_build_node = ModelNode(
    label="AutoBuild",
    model=SequentialMLP(output_shape=(1, 1), n_layers=1, hidden_dim=16),
    upstream_ref=fs_ref,
)
print(f"Before forward: is_built={auto_build_node.is_built}")

# First forward pass triggers auto-build
with torch.no_grad():
    output = auto_build_node(sample_data)

print(f"After forward:  is_built={auto_build_node.is_built}")
print(f"  input_shape:  {auto_build_node.input_shape}")
print(f"  output_shape: {auto_build_node.output_shape}")

We could’ve omitted output_shape, which results in the same (1,1) shape (because the FeatureSet 'soh' data has shape (1,1)).

It is generally best practice to explicitly define the output shape of any models you create.

---

Chaining Nodes

A ModelNode can take another ModelNode (or any ComputeNode) as its upstream, enabling multi-stage pipelines.

# Encoder -> Regressor chain
encoder = ModelNode(
    label="Encoder",
    model=SequentialMLP(output_shape=(1, 8), n_layers=2, hidden_dim=32),
    upstream_ref=fs_ref,
)

regressor = ModelNode(
    label="Regressor",
    model=SequentialMLP(output_shape=(1, 1), n_layers=1, hidden_dim=16),
    upstream_ref=encoder,  # Receives output from Encoder
)

print(f"Encoder upstream:   {encoder.upstream_ref.resolve()}")
print(f"Regressor upstream: {regressor.upstream_ref.resolve()}")

---

Freezing and Unfreezing

Freezing a node prevents its parameters from being updated during training. This is useful for transfer learning or multi-stage training events.

The below requires_grad property is PyTorch-specific, but similar gradient blocking is enforced for TensorFlow models.

print(f"Frozen: {node_mlp.is_frozen}")

node_mlp.freeze()
print(f"After freeze:   {node_mlp.is_frozen}")

# Verify PyTorch parameters are frozen
param = next(node_mlp.model.parameters())
print(f" - requires_grad:  {param.requires_grad}")

node_mlp.unfreeze()
print(f"After unfreeze: {node_mlp.is_frozen}")
print(f" - requires_grad:  {param.requires_grad}")

---

Serialization

ModelNode supports full config and state serialization via get_config() / from_config() and get_state() / set_state(). The underlying BaseModel handles weight serialization.

# Configuration (structure, no weights)
config = node_mlp.get_config()
print("Config keys:", list(config.keys()))

# State (includes learned weights)
state = node_mlp.get_state()
print("State keys:", list(state.keys()))
print("Model weight keys:", list(state["model"]["weights"].keys())[:3], "...")

Models can also be saved to and loaded from disk independently of the ModelNode.

Note that save and load methods are not provided on ModelNodes, only on the BaseModel itself. This is intentional. ModelNodes are not useful outside of its parent Experiment (their upstream and downstream connections have no meaning on their own). However, the underlying model is useful to share independently.

from pathlib import Path
from tempfile import TemporaryDirectory

from modularml import BaseModel

SAVE_DIR = TemporaryDirectory()

# Save and reload a built-in model
save_path = node_mlp.model.save(Path(SAVE_DIR.name) / "my_mlp", overwrite=True)
reloaded = BaseModel.load(save_path)
print(f"Models equal: {reloaded == node_mlp.model}")

For custom (non-built-in) models, the model’s source code is packaged alongside the weights. This requires that the custom model be defined in a standalone python file.

# Save a wrapped custom model
node_custom.build_model(input_shape=(10,), output_shape=(4,))
try:
    save_path_custom = node_custom.model.save(
        Path(SAVE_DIR.name) / "my_custom",
        overwrite=True,
    )
except RuntimeError as e:
    print(e)

from utils.my_model import MyEncoder

# After moving the MyEncoder class to an external python file, we can save
custom_node = ModelNode(
    label="imported_model",
    model=MyEncoder(in_features=10, out_features=4),
    upstream_ref=fs_ref,
)

save_path_custom = custom_node.model.save(
    Path(SAVE_DIR.name) / "my_custom",
    overwrite=True,
)

Packaging source code is the only way to ensure full reproducibility of custom code. However, you should always inspect unknown code below executing it.

If you try to load a saved file that contains packaged code, an error will be thrown unless you intentionally set allow_packaged_code=True.

The following procedure is recommended when loading unknown serialized mml files:

1. Call load(..., allow_packaged_code=False) (it will always defaul to False)

2. If an error occurs, indicating packaged code, use the following utility to inspect the source code before importing as an executable: modularml.utils.inspect_packaged_code

3. After verifying the nature of the source code, you can retry load with allow_packaged_code=True

# Reload with packaged code
try:
    reloaded_custom = BaseModel.load(save_path_custom, allow_packaged_code=False)
except RuntimeError as e:
    print(e)

from modularml.utils import inspect_packaged_code

# Inspect the file before reloadinge
# The method returns a dict with keys for each class that would need to be loaded
res = inspect_packaged_code(save_path_custom)
for k, code in res.items():
    print(k)
    print(code)

# Now that we've verified the code is safe to run, we can try loading again
reloaded_custom = BaseModel.load(save_path_custom, allow_packaged_code=True)
print(f"Models equal: {reloaded_custom == custom_node.model}")

---

Summary

Model Classes

Class	Module	Backend	Description
BaseModel	modularml.core.models	Abstract	Base interface for all models.
TorchBaseModel	modularml.core.models	PyTorch	Base for built-in PyTorch models.
TorchModelWrapper	modularml.core.models	PyTorch	Wraps any torch.nn.Module.
TensorflowModelWrapper	modularml.core.models	TensorFlow	Wraps any tf.keras.Model.
ScikitModelWrapper	modularml.core.models	scikit-learn	Wraps any sklearn.BaseEstimator.
SequentialMLP	modularml.models.torch	PyTorch	Built-in multi-layer perceptron.
SequentialCNN	modularml.models.torch	PyTorch	Built-in 1D convolutional network.

ModelNode Properties and Methods

Property / Method	Description
.model	The underlying BaseModel instance.
.backend	Backend enum (Backend.TORCH, Backend.TENSORFLOW, etc.).
.is_built	Whether the model has been built with input/output shapes.
.input_shape	Input shape tuple (no batch dim), or None.
.output_shape	Output shape tuple (no batch dim), or None.
.upstream_ref	The single upstream reference (read-only property).
.is_frozen	Whether training is disabled for this node.
build_model(input_shape, output_shape)	Build the model and optimizer manually.
forward_single(x) / __call__(x)	Forward pass on SampleData, RoleData, or Batch.
freeze() / unfreeze()	Toggle parameter trainability.
get_config() / from_config()	Config serialization (structure only).
get_state() / set_state()	State serialization (includes weights).

Next Steps

• ModelGraph: Compose multiple ModelNodes (and MergeNodes) into a computational graph that handles build order, shape inference, and forward pass routing.

• Experiment: Use Experiment to combine a ModelGraph with training phases, loss functions, and evaluation — the primary user-facing entry point.