class

PretrainedModel

extendsModule

PretrainedModel(config: ModelConfig)

source edit

Base for every Lucid model that supports the pretrained-checkpoint flow.

Subclasses inherit from_pretrained / save_pretrained plus parameter-counting and embedding-access helpers. The contract:

Set config_class = MyConfig as a ClassVar — required; enforced at __init__ time.
Define __init__(self, config: MyConfig) -> None taking a single config argument. All architectural variation (depth, width, activation choices, …) belongs inside the config — no extra constructor parameters.
Implement forward(...) -> ModelOutput returning one of the dataclasses in lucid.models._output.

Attributes

config_classClassVar[type[ModelConfig] or None]

Subclasses MUST override. None signals "not set" — the base __init__ raises in that case.

base_model_prefixClassVar[str]

Optional name of the backbone attribute (e.g. "bert" on BERTForMaskedLM). Used by future state-dict remapping logic when loading head-less checkpoints into head-bearing models.

configModelConfig

Instance attribute populated by __init__.

Notes

The class is a thin lift on lucid.nn.Module — all parameter management still flows through the standard Module machinery. Two persistence formats are supported when saving / loading: the pickle-based weights.lucid (default) and the cross-framework model.safetensors.

Examples

>>> class MyModel(PretrainedModel):
...     config_class: ClassVar[type[MyConfig]] = MyConfig
...     def __init__(self, config):
...         super().__init__(config)
...         self.linear = nn.Linear(config.hidden_size, config.num_classes)
...     def forward(self, x):
...         return self.linear(x)
>>> model = MyModel(MyConfig(hidden_size=128, num_classes=10))
>>> model.num_parameters()
1290

Used by 55

… 43 more

Constructors

dunder

init

→None

__init__(config: ModelConfig)

source edit

Initialise the module and validate the supplied config.

Parameters

configModelConfig

Must be an instance of the subclass's declared config_class.

Raises

TypeError

If config_class has not been set on the concrete subclass, or if config is not an instance of config_class.

Class methods

cls

from_pretrained

→Self

from_pretrained(name_or_path: str, strict: bool = True)

source edit

Load a model from a registered name or a local directory.

Two modes are supported:

Registered name ("resnet_50" / "resnet-50") — looked up in the global registry; the factory is invoked with pretrained=True. The factory result is validated to be an instance of cls (so subclass calls like ResNet.from_pretrained("vit_base_16") raise rather than silently returning the wrong family).
Local directory containing config.json plus either model.safetensors (preferred) or weights.lucid. The config is restored via cls.config_class.load and the weights are loaded with lucid.load.

Parameters

name_or_pathstr

Registered model name or filesystem directory path.

strict(bool, optional, keyword - only)= True

Forwarded to load_state_dict.

Returns

Self

A fully constructed model instance.

Raises

ValueError

If the name is unrecognised and is not a valid directory.

FileNotFoundError

If a directory was supplied but required files are absent.

TypeError

If the registered factory yields a non-subclass of cls, or config_class is unset, or the weights file lacks a state-dict.

Notes

For task-aware dispatch that resolves to a different concrete subclass per task, use the AutoModelFor* family instead.

Examples

>>> # Registered name
>>> model = ResNetForImageClassification.from_pretrained("resnet_50")
>>>
>>> # Local directory
>>> model.save_pretrained("/tmp/my_resnet50")
>>> reloaded = ResNetForImageClassification.from_pretrained("/tmp/my_resnet50")

Instance methods

get_input_embeddings

→nn.Module or None

get_input_embeddings()

source edit

Return the input-embedding submodule, or None if not applicable.

Returns

nn.Module or None

The embedding layer for text / token-id models (BERT, GPT, …); None for vision and other non-embedding models.

Notes

Override in subclasses that have an embedding table. Used by tools that need to resize / share / introspect token embeddings without coupling to family-specific attribute names.

num_parameters

→int

num_parameters(only_trainable: bool = False)

source edit

Return the total number of elements across all parameters.

Parameters

only_trainable(bool, optional, keyword - only)= False

When True, parameters with requires_grad=False are excluded (useful for reporting trainable model size after freezing the backbone).

Returns

int

Sum of prod(p.shape) over the selected parameters.

Examples

>>> model = create_model("resnet_50")
>>> model.num_parameters()
25557032
>>> for p in model.backbone.parameters():
...     p.requires_grad = False
>>> model.num_parameters(only_trainable=True) < 25557032
True

save_pretrained

→None

save_pretrained(path: str, safe_serialization: bool = False)

source edit

Write config.json and weights to path.

Parameters

pathstr

Destination directory. Created if it does not exist.

safe_serialization(bool, optional, keyword - only)= False

If True, save weights as model.safetensors (requires pip install safetensors). If False, save as weights.lucid using the native pickle-based format.

Notes

Output layout:

path/
  config.json
  model.safetensors   # when safe_serialization=True
  weights.lucid       # when safe_serialization=False

The companion from_pretrained (or any AutoModelFor* class) reads this directory layout and prefers SafeTensors when both files are present.

Examples

>>> model = create_model("resnet_50")
>>> model.save_pretrained("/tmp/my_resnet50", safe_serialization=True)

set_input_embeddings

→None

set_input_embeddings(value: nn.Module)

source edit

Replace the input-embedding submodule.

Parameters

valuenn.Module

New embedding module.

Raises

NotImplementedError

On the base class — subclasses with embeddings must override.

Notes

Companion to get_input_embeddings. Subclasses that override one should override the other.

>>> class MyModel(PretrainedModel): ... config_class: ClassVar[type[MyConfig]] = MyConfig ... def __init__(self, config): ... super().__init__(config) ... self.linear = nn.Linear(config.hidden_size, config.num_classes) ... def forward(self, x): ... return self.linear(x) >>> model = MyModel(MyConfig(hidden_size=128, num_classes=10)) >>> model.num_parameters() 1290

>>> # Registered name >>> model = ResNetForImageClassification.from_pretrained("resnet_50") >>> >>> # Local directory >>> model.save_pretrained("/tmp/my_resnet50") >>> reloaded = ResNetForImageClassification.from_pretrained("/tmp/my_resnet50")

>>> model = create_model("resnet_50") >>> model.num_parameters() 25557032 >>> for p in model.backbone.parameters(): ... p.requires_grad = False >>> model.num_parameters(only_trainable=True) < 25557032 True

save_pretrained

→None

save_pretrained(path: str, safe_serialization: bool = False)

source edit

Write config.json and weights to path.

Parameters

pathstr

Destination directory. Created if it does not exist.

safe_serialization(bool, optional, keyword - only)= False

If True, save weights as model.safetensors (requires pip install safetensors). If False, save as weights.lucid using the native pickle-based format.

Notes

Output layout:

path/
  config.json
  model.safetensors   # when safe_serialization=True
  weights.lucid       # when safe_serialization=False

The companion from_pretrained (or any AutoModelFor* class) reads this directory layout and prefers SafeTensors when both files are present.

Examples

>>> model = create_model("resnet_50")
>>> model.save_pretrained("/tmp/my_resnet50", safe_serialization=True)