GitHub

Porting

This page is for contributors, and for a package that registers architectures of its own. It is not public API: the names below are neither exported nor declared public, and they may change without notice, in any release.

Metalhead's libraries (timm, transformers, diffusers and sentence-transformers) plug into create_model through a registry. An AbstractConfigFormat says how the config.json of a repository names an architecture and its arguments, an ArchSpec is a ported architecture in one format, and register! adds it to the registry that create_model and list_models read. src/Timm/config.jl is the template of a format, and Porting a model the workflow. The processors have their own interface (register_processor! and the others), listed in the Internals.

using Metalhead, Flux
using Metalhead: ArchSpec, AbstractConfigFormat, register!

struct ToyFormat <: AbstractConfigFormat end
Metalhead.architecture_of(::ToyFormat, config) = get(config, "toy_architecture", nothing)

toy_mlp(; in = 3, out = 2) = (; fc = Dense(in => out))   # weight keys follow the field names
register!(ArchSpec(ToyFormat(), "toy_mlp", toy_mlp; family = "toy"))

create_model("toy_mlp"; in = 5)             # a model with initial weights
create_model("path/to/dir"; pretrained = true)  # config.json and model.safetensors in it

Architectures

Metalhead.ArchSpec — Type
ArchSpec(format, name, build; family = "other", defaults = (;),
         references = Pair{String,String}[], ignore = Regex[], ignore_config = String[])

A ported architecture, as config.json files in format name it. build(; kwargs...) constructs the model from the keyword arguments of the original library (e.g. timm's), merged over defaults. references lists repo => sha checkpoints the port was validated against; the first one is the default pretrained checkpoint. ignore lists extra checkpoint keys the model doesn't use, and ignore_config the keys of config.json that it doesn't read on purpose (training-only settings, ...) on top of those that format ignores (ignores_config_key): create_model warns about every other key that build doesn't declare. family is the original library's family identifier, used to group architectures in list_models and the documentation.

source

Config formats

A format is a subtype of AbstractConfigFormat with methods of these functions.

Metalhead.AbstractConfigFormat — Type
AbstractConfigFormat

How the config.json of a model repository names an architecture and its arguments: timm's (architecture, model_args, pretrained_cfg), a future library's architectures, ... A ported library defines a subtype and registers its architectures with it. Only architecture_of is required; the other methods have defaults.

source
Metalhead.architecture_of — Function
architecture_of(format, config) -> Union{Nothing, String}

The name of the architecture described by the parsed config.json, or nothing if the configuration isn't in this format.

source
Metalhead.model_args — Function
model_args(format, spec, config, kwargs::NamedTuple) -> NamedTuple

The keyword arguments for spec.build: the defaults of the architecture, overridden by the configuration and then by the user's kwargs.

source
Metalhead.ignores_config_key — Function
ignores_config_key(format, key) -> Bool

Whether every architecture of format ignores the key key of a configuration on purpose: what its library records about the checkpoint, or sets for training (transformers' architectures, transformers_version, ...). Other keys that a model doesn't read warn when it is created, and an architecture lists its own in ArchSpec's ignore_config. By default no key is ignored.

source
Metalhead.skip_prefix — Function
skip_prefix(format, config, args) -> Union{Nothing, String}

The prefix of the checkpoint keys that don't load into a model built with args, which keeps its initial values (e.g. the classifier after changing num_classes). Throws if the checkpoint can't be loaded into such a model at all.

source
Metalhead.weights_name — Function
weights_name(format) -> String

The name of the safetensors weights of a model in format: "model" for model.safetensors or its shards (the default).

source
Metalhead.instantiate — Function
instantiate(format, spec, source, config; pretrained, eltype, kwargs...)

Build the model spec of the parsed config, from model_args, in the element type eltype, and with pretrained = true load its weights from source. A format whose configuration describes several models (a diffusers pipeline) builds them here.

source

Checkpoints

Metalhead.Checkpoint — Type
Checkpoint(paths)

Lazy view over one or more safetensors files (e.g. the shards of a sharded checkpoint). Indexing with a PyTorch state_dict key returns the tensor in Julia layout, i.e. with the PyTorch dimensions reversed (NCHW → WHCN).

source
Metalhead.load_state! — Function
load_state!(model, ckpt::Checkpoint; prefix = "", ignore = [r"num_batches_tracked$"],
            skip = nothing, strict = true, filter_fn = nothing) -> NamedTuple

Copy the tensors of ckpt into model, matching PyTorch state_dict keys derived from the model structure (see torch_params) and nested under prefix. Keys matching a regex in ignore may be absent from the model; parameters whose key starts with skip keep their current values. With strict = true, missing keys, unused checkpoint keys and shape mismatches throw. filter_fn(param, x) (see checkpoint_tensor) replaces each tensor x read for the TorchParam param, in the layout of its array, before its shape is checked. Returns the number of tensors loaded and the lists of missing, unexpected and mismatched keys.

source