state_dict

Checkpoint reading, adaptive state_dict loading, and key-mapping transforms.

Functions:

read_state_dict –

Read a checkpoint file into a flat state dict.
load_state_dict –

Load state_dict into model, keeping the initialization of the keys whose shapes do not match.
transform_state_dict –

Transform a pytorch module state dict by remapping keys and optionally transforming associated tensors.

read_state_dict [source] ¶

read_state_dict(path: PathLike) -> Dict[str, Any]

Read a checkpoint file into a flat state dict.

Supports .safetensors and torch.save files, unwrapping a state_dict or model_state_dict key if present. Lightning checkpoints (identified by their pytorch-lightning_version key) are reduced to the wrapped network: only model.-prefixed keys are kept, with the prefix stripped.

Parameters:

path (PathLike) –

Checkpoint file to read.

Returns:

Dict[str, Any] –

The flat parameter-name to tensor mapping.

Raises:

ValueError –

If the checkpoint is not a mapping, or if a Lightning checkpoint has no model.-prefixed key to extract.

Example

from torch_pointcloud.utils.state_dict import read_state_dict

state_dict = read_state_dict("checkpoints/last.ckpt")

load_state_dict [source] ¶

load_state_dict(
    model: Module,
    state_dict: Dict[str, Any],
    source: str,
    strict: bool = True,
) -> None

Load state_dict into model, keeping the initialization of the keys whose shapes do not match.

Checkpoint keys absent from the model are ignored with a warning, keys with mismatched tensor shapes are skipped and keep their fresh initialization (e.g. the stem after an in_channels override), and model keys missing from the checkpoint raise. A checkpoint matching the model exactly loads completely, as with a strict load.

Parameters:

model (Module) –

Module to load the parameters into.
state_dict (Dict[str, Any]) –

Flat parameter-name to tensor mapping (see read_state_dict).
source (str) –

Label identifying where the checkpoint came from, used in warnings and errors.
strict (bool, default: True ) –

Raise when the checkpoint is missing model keys. Pass False to warm-start from a partial checkpoint (e.g. a backbone-only or pretraining export): missing keys warn and keep their fresh initialization.

Raises:

RuntimeError –

If strict and the checkpoint is missing keys the model requires.

Example

from torch_pointcloud.utils.state_dict import load_state_dict, read_state_dict

load_state_dict(model, read_state_dict("checkpoints/last.ckpt"), source="last.ckpt")

transform_state_dict [source] ¶

transform_state_dict(
    state_dict: Dict[str, Any],
    mapping: Dict[str, str],
    value_transform: Optional[
        Callable[[Tensor], Tensor]
    ] = None,
    strict: bool = False,
) -> Dict[str, Any]

Transform a pytorch module state dict by remapping keys and optionally transforming associated tensors. This function is designed to map the state dict of a pytorch module to a different state dict, facilitating the transfer of weights between different models.

Parameters:

state_dict (Dict[str, Any]) –

The state dict to transform.
mapping (Dict[str, str]) –

A dictionary mapping the old keys to the new keys.
value_transform (Optional[Callable[[Tensor], Tensor]], default: None ) –

A function to transform the values.
strict (bool, default: False ) –

Raise a ValueError on mapping patterns that match no key and on source keys colliding onto the same destination key. When False, unused patterns are ignored and collisions only emit a warning.

Returns:

Dict[str, Any] –

The transformed state dict.

Example

import torch
from torch_pointcloud.utils.state_dict import transform_state_dict

state_dict = {
    "encoder.conv.0.weight": torch.randn(1, 3, 16, 16),
    "encoder.conv.0.bias": torch.randn(1),
    "encoder.norm.1.weight": torch.randn(1, 16, 16, 16),
    "encoder.norm.1.bias": torch.randn(1),
    "encoder.norm.1.running_mean": torch.randn(16),
    "encoder.norm.1.running_var": torch.randn(16),
}
mapping = {
    "encoder.{module}.{i}.weight": "backbone.{module}.{i+1}.weight",
    "encoder.{module}.{i}.bias": "backbone.{module}.{i+1}.bias",
    "encoder.{module}.{i}.running_{stat}": "backbone.{module}.{i+1}.running_{stat}",
}
state_dict = transform_state_dict(state_dict, mapping)
print(state_dict.keys())
# {
#     "backbone.conv.1.weight": ...,
#     "backbone.conv.1.bias": ...,
#     "backbone.norm.2.weight": ...,
#     "backbone.norm.2.bias": ...,
#     "backbone.norm.2.running_mean": ...,
#     "backbone.norm.2.running_var": ...,
# }