I/O & Export (`pypielm.io`)¶

I/O: checkpointing and model export.

Public surface:

from pypielm.io import (
    save_model, load_model,
    to_onnx, to_torchscript,
)

pypielm.io.save_model(model, path, *, include_config=True, overwrite=False)[source]¶

Serialise model weights (and optionally config) to path.

The checkpoint format is a torch.save-compatible dict:

{
    "version": "0.1.0",
    "model_class": "<registry name or qualified class name>",
    "state_dict": { ... },
    "config": { ... },   # only when include_config=True
}

Parameters:

model (BasePIELM) – A fitted BasePIELM instance.
path (str | Path) – Destination file path (.pt extension recommended).
include_config (bool) – Whether to embed the model’s config in the checkpoint.
overwrite (bool) – If False, raise FileExistsError when path already exists.

Return type:

None

pypielm.io.load_model(path, *, model_class=None, device='cpu', dtype=None)[source]¶

Load a checkpoint written by save_model().

If model_class is None, the class is inferred from the checkpoint’s model_class field via the model registry.

Parameters:

path (str | Path) – Path to the checkpoint file.
model_class (type[BasePIELM] | None) – Override the class to use for reconstruction.
device (str | device) – Device to load the model onto.
dtype (dtype | None) – Dtype override (default: use saved dtype).

Return type:

BasePIELM

Returns:

A BasePIELM instance with weights loaded.

pypielm.io.to_onnx(model, path, *, example_input=None, input_dim=2, opset_version=17)[source]¶

Export model to ONNX format.

Requires onnx and onnxruntime (install via pip install pypielm[export]).

Parameters:

model (BasePIELM) – A fitted PIELM model.
path (str | Path) – Destination .onnx file path.
example_input (Tensor | None) – Representative input tensor (1, d) for tracing. If None, a zero tensor of shape (1, input_dim) is used.
input_dim (int) – Spatial dimension d (used when example_input is None).
opset_version (int) – ONNX opset version.

Return type:

None

pypielm.io.to_torchscript(model, path, *, example_input=None, input_dim=2, method='trace')[source]¶

Export model to TorchScript.

Parameters:

model (BasePIELM) – A fitted PIELM model.
path (str | Path) – Destination .pt file path.
example_input (Tensor | None) – Representative input for tracing.
input_dim (int) – Spatial dimension (used when example_input is None).
method (str) – 'trace' or 'script'.

Return type:

ScriptModule

Returns:

The compiled torch.jit.ScriptModule.

Checkpointing¶

Model checkpointing: save and load trained PIELM weights.

pypielm.io.checkpoint.save_model(model, path, *, include_config=True, overwrite=False)[source]¶

Serialise model weights (and optionally config) to path.

The checkpoint format is a torch.save-compatible dict:

{
    "version": "0.1.0",
    "model_class": "<registry name or qualified class name>",
    "state_dict": { ... },
    "config": { ... },   # only when include_config=True
}

Parameters:

model (BasePIELM) – A fitted BasePIELM instance.
path (str | Path) – Destination file path (.pt extension recommended).
include_config (bool) – Whether to embed the model’s config in the checkpoint.
overwrite (bool) – If False, raise FileExistsError when path already exists.

Return type:

None

pypielm.io.checkpoint.load_model(path, *, model_class=None, device='cpu', dtype=None)[source]¶

Load a checkpoint written by save_model().

If model_class is None, the class is inferred from the checkpoint’s model_class field via the model registry.

Parameters:

path (str | Path) – Path to the checkpoint file.
model_class (type[BasePIELM] | None) – Override the class to use for reconstruction.
device (str | device) – Device to load the model onto.
dtype (dtype | None) – Dtype override (default: use saved dtype).

Return type:

BasePIELM

Returns:

A BasePIELM instance with weights loaded.

Model Export¶

Model export to portable inference formats.

pypielm.io.export.to_onnx(model, path, *, example_input=None, input_dim=2, opset_version=17)[source]¶

Export model to ONNX format.

Requires onnx and onnxruntime (install via pip install pypielm[export]).

Parameters:

model (BasePIELM) – A fitted PIELM model.
path (str | Path) – Destination .onnx file path.
example_input (Tensor | None) – Representative input tensor (1, d) for tracing. If None, a zero tensor of shape (1, input_dim) is used.
input_dim (int) – Spatial dimension d (used when example_input is None).
opset_version (int) – ONNX opset version.

Return type:

None

pypielm.io.export.to_torchscript(model, path, *, example_input=None, input_dim=2, method='trace')[source]¶

Export model to TorchScript.

Parameters:

model (BasePIELM) – A fitted PIELM model.
path (str | Path) – Destination .pt file path.
example_input (Tensor | None) – Representative input for tracing.
input_dim (int) – Spatial dimension (used when example_input is None).
method (str) – 'trace' or 'script'.

Return type:

ScriptModule

Returns:

The compiled torch.jit.ScriptModule.

I/O & Export (pypielm.io)¶

Checkpointing¶

Model Export¶

I/O & Export (`pypielm.io`)¶