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

>>> import lucid
>>> import lucid.nn as nn
>>> from dataclasses import dataclass
>>> from typing import ClassVar
>>> from lucid.models import ModelConfig, PretrainedModel
>>> @dataclass(frozen=True)
... class MyConfig(ModelConfig):
...     model_type: ClassVar[str] = "myfamily"
...     hidden_size: int = 768
...     num_classes: int = 10
>>> 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 58

… 46 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

Trunks with an embedding table override this. A task wrapper does not need to: this implementation hands the call to the trunk stored under base_model_prefix, so a wrapper reports the same table as the trunk inside it. Used by tools that need to resize / share / introspect token embeddings without coupling to family-specific attribute names.

Examples

>>> from lucid.models import create_model
>>> lm = create_model("gpt_lm", vocab_size=100, hidden_size=16,
...                   num_hidden_layers=1, num_attention_heads=2,
...                   intermediate_size=32, max_position_embeddings=8)
>>> lm.get_input_embeddings() is lm.transformer.tokens_embed
True

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

When the model has no embedding table to replace — neither an override of its own nor a trunk under base_model_prefix.

Notes

Companion to get_input_embeddings. Trunks with a table override both; a task wrapper inherits this implementation, which makes the swap on its trunk.

When the wrapper ties an output head to the table (config.tie_word_embeddings), the head is re-bound to the new table's weight, so the two stay one matrix and the head scores the new vocabulary. A head bias sized to the vocabulary is resized with it: existing entries are kept and new ones start at zero. An untied head is left as it is. Either way config is frozen and keeps its original vocab_size.

Examples

A text trunk hands its table over, so a larger one can be swapped
in — here to make room for twenty token ids the original 100-row
table could not look up.
>>> import lucid
>>> import lucid.nn as nn
>>> from lucid.models import create_model
>>> model = create_model("gpt", vocab_size=100, hidden_size=16,
...                      num_hidden_layers=1, num_attention_heads=2,
...                      intermediate_size=32,
...                      max_position_embeddings=8).eval()
>>> bigger = nn.Embedding(120, 16)
>>> model.set_input_embeddings(bigger)
>>> model.get_input_embeddings() is bigger
True
>>> ids = lucid.tensor([[110, 3]], dtype=lucid.int64)
>>> model(ids).last_hidden_state.shape
(1, 2, 16)
A task wrapper makes the same swap on its trunk.  Its LM head is tied
to the table, so the head follows and the logits widen to the new
vocabulary, rather than keep scoring the table it replaced.
>>> lm = create_model("gpt_lm", vocab_size=100, hidden_size=16,
...                   num_hidden_layers=1, num_attention_heads=2,
...                   intermediate_size=32,
...                   max_position_embeddings=8).eval()
>>> lm.set_input_embeddings(nn.Embedding(120, 16))
>>> lm.lm_head.weight is lm.get_input_embeddings().weight
True
>>> lm(ids).logits.shape
(1, 2, 120)
A model with no embedding table keeps this base implementation and
refuses, rather than attach a module that nothing would read.
>>> lenet = create_model("lenet_5")
>>> lenet.get_input_embeddings() is None
True
>>> lenet.set_input_embeddings(nn.Embedding(10, 4))
Traceback (most recent call last):
    ...
NotImplementedError: LeNet does not support set_input_embeddings

>>> import lucid >>> import lucid.nn as nn >>> from dataclasses import dataclass >>> from typing import ClassVar >>> from lucid.models import ModelConfig, PretrainedModel >>> @dataclass(frozen=True) ... class MyConfig(ModelConfig): ... model_type: ClassVar[str] = "myfamily" ... hidden_size: int = 768 ... num_classes: int = 10 >>> 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")

>>> from lucid.models import create_model >>> lm = create_model("gpt_lm", vocab_size=100, hidden_size=16, ... num_hidden_layers=1, num_attention_heads=2, ... intermediate_size=32, max_position_embeddings=8) >>> lm.get_input_embeddings() is lm.transformer.tokens_embed True

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

A text trunk hands its table over, so a larger one can be swapped in — here to make room for twenty token ids the original 100-row table could not look up. >>> import lucid >>> import lucid.nn as nn >>> from lucid.models import create_model >>> model = create_model("gpt", vocab_size=100, hidden_size=16, ... num_hidden_layers=1, num_attention_heads=2, ... intermediate_size=32, ... max_position_embeddings=8).eval() >>> bigger = nn.Embedding(120, 16) >>> model.set_input_embeddings(bigger) >>> model.get_input_embeddings() is bigger True >>> ids = lucid.tensor([[110, 3]], dtype=lucid.int64) >>> model(ids).last_hidden_state.shape (1, 2, 16) A task wrapper makes the same swap on its trunk. Its LM head is tied to the table, so the head follows and the logits widen to the new vocabulary, rather than keep scoring the table it replaced. >>> lm = create_model("gpt_lm", vocab_size=100, hidden_size=16, ... num_hidden_layers=1, num_attention_heads=2, ... intermediate_size=32, ... max_position_embeddings=8).eval() >>> lm.set_input_embeddings(nn.Embedding(120, 16)) >>> lm.lm_head.weight is lm.get_input_embeddings().weight True >>> lm(ids).logits.shape (1, 2, 120) A model with no embedding table keeps this base implementation and refuses, rather than attach a module that nothing would read. >>> lenet = create_model("lenet_5") >>> lenet.get_input_embeddings() is None True >>> lenet.set_input_embeddings(nn.Embedding(10, 4)) Traceback (most recent call last): ... NotImplementedError: LeNet does not support set_input_embeddings