Skip to content

deeplabcut.core.weight_init

Classes to configure how to initialize model weights.

Classes:

Name Description
WeightInitialization

Configures weights initialization when transfer learning or fine-tuning models.

WeightInitialization

Bases: DLCBaseConfig

Configures weights initialization when transfer learning or fine-tuning models.

Parameters:

Name Type Description Default

snapshot_path

The path to the snapshot used to initialize pose model weights when training a model.

 required

detector_snapshot_path

The path to the snapshot used to initialize detector weights when training a model.

 required

dataset

Optionally, the dataset on which the snapshots were trained. Required when fine-tuning SuperAnimal models.

 required

with_decoder

Whether to load the decoder weights as well.

 required

memory_replay

Only when with_decoder=True. Whether to train the model with memory replay, so that it predicts all SuperAnimal (or previous project) bodyparts.

 required

conversion_array

The mapping from SuperAnimal (or other project, on which the weights were trained) to project bodyparts. Required when with_decoder=True. An array [7, 0, 1] means the project has 3 bodyparts, where the 1st bodypart corresponds to the 8th bodypart in the pretrained model, the 2nd to the 1st and the 3rd to the 2nd (as arrays are 0-indexed).

 required

bodyparts

Optionally, the name of each bodypart entry in the conversion array.

 required

Methods:

Name Description
build

Builds a WeightInitialization for a project.
from_dict_legacy

Deals with weight initialization that were created before 3.0.0rc5
Source code in deeplabcut/core/weight_init.py 
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
class WeightInitialization(DLCBaseConfig):
    """Configures weights initialization when transfer learning or fine-tuning models.

    Args:
        snapshot_path: The path to the snapshot used to initialize pose model weights
            when training a model.
        detector_snapshot_path: The path to the snapshot used to initialize detector
            weights when training a model.
        dataset: Optionally, the dataset on which the snapshots were trained. Required
            when fine-tuning SuperAnimal models.
        with_decoder: Whether to load the decoder weights as well.
        memory_replay: Only when ``with_decoder=True``. Whether to train the model with
            memory replay, so that it predicts all SuperAnimal (or previous project)
            bodyparts.
        conversion_array: The mapping from SuperAnimal (or other project, on which the
            weights were trained) to project bodyparts. Required when
            `with_decoder=True`.
            An array [7, 0, 1] means the project has 3 bodyparts, where the 1st bodypart
            corresponds to the 8th bodypart in the pretrained model, the 2nd to the 1st
            and the 3rd to the 2nd (as arrays are 0-indexed).
        bodyparts: Optionally, the name of each bodypart entry in the conversion array.
    """

    snapshot_path: Path | None = None
    detector_snapshot_path: Path | None = None
    dataset: str | None = None
    with_decoder: bool = False
    memory_replay: bool = False
    conversion_array: NDArrayInt | None = None
    bodyparts: list[str] | None = None

    @model_validator(mode="after")
    def _validate_options(self) -> Self:
        if self.memory_replay and not self.with_decoder:
            raise ValueError(
                "You cannot train a model with memory replay if you do not keep the "
                "decoder layers (``with_decoder=True``), but you passed "
                "`memory_replay=True` and `with_decoder=False`. Please change your "
                "WeightInitialization parameters."
            )

        if self.with_decoder and self.conversion_array is None:
            raise ValueError(
                "You must specify a conversion_array to initialize decoder weights (``with_decoder=True``)."
            )

        if self.bodyparts is not None and self.conversion_array is None:
            raise ValueError(
                "Specifying bodyparts should only be done when `with_decoder=True` and"
                " the conversion array is specified."
            )

        if self.conversion_array is not None and self.bodyparts is not None:
            if len(self.conversion_array) != len(self.bodyparts):
                raise ValueError(
                    f"There must be the same number of elements in the bodyparts list "
                    f"and conv. array; found {self.bodyparts}, {self.conversion_array}"
                )
        return self

    @field_validator("snapshot_path", "detector_snapshot_path", mode="before")
    @classmethod
    def _coerce_null_path(cls, v):
        if v is None or v == "None":
            return None
        return v

    @classmethod
    def from_dict(cls, data: dict) -> Self:
        if "snapshot_path" not in data:
            return cls.from_dict_legacy(data)
        return cls.model_validate(data)

    @classmethod
    def from_dict_legacy(cls, data: dict) -> Self:
        """Deals with weight initialization that were created before 3.0.0rc5"""

        import deeplabcut.pose_estimation_pytorch.modelzoo.utils as utils

        conversion_array = data.get("conversion_array")
        if conversion_array is not None:
            conversion_array = np.array(conversion_array, dtype=int)

        return cls(
            snapshot_path=utils.get_super_animal_snapshot_path(
                dataset=data["dataset"],
                model_name="hrnet_w32",
            ),
            detector_snapshot_path=utils.get_super_animal_snapshot_path(
                dataset=data["dataset"],
                model_name="fasterrcnn_resnet50_fpn_v2",
            ),
            with_decoder=data["with_decoder"],
            memory_replay=data["memory_replay"],
            conversion_array=conversion_array,
            bodyparts=data.get("bodyparts"),
        )

    @staticmethod
    def build(
        cfg: dict,
        super_animal: str,
        model_name: str = "hrnet_w32",
        detector_name: str = "fasterrcnn_resnet50_fpn_v2",
        with_decoder: bool = False,
        memory_replay: bool = False,
        customized_pose_checkpoint: str | None = None,
        customized_detector_checkpoint: str | None = None,
    ) -> WeightInitialization:
        """Builds a WeightInitialization for a project.

        `WeightInitialization.build` is deprecated and will be removed in a future
        version of DeepLabCut. Please use `build_weight_init` from `deeplabcut.modelzoo`
        instead.

        Args:
            cfg: The project's configuration.
            super_animal: The SuperAnimal model with which to initialize weights.
            model_name: The name of the model architecture for which to load the weights
                (defaults to "hrnet_w32" for backwards compatibility).
            detector_name: The name of the detector architecture for which to load the
                weights (defaults to "fasterrcnn_resnet50_fpn_v2" for backwards
                compatibility).
            with_decoder: Whether to load the decoder weights as well. If this is true,
                a conversion table must be specified for the given SuperAnimal in the
                project configuration file. See
                ``deeplabcut.modelzoo.utils.create_conversion_table`` to create a
                conversion table.
            memory_replay: Only when ``with_decoder=True``. Whether to train the model
                with memory replay, so that it predicts all SuperAnimal bodyparts.
            customized_pose_checkpoint: A customized SuperAnimal pose checkpoint, as an
                alternative to the Hugging Face one
            customized_detector_checkpoint: A customized SuperAnimal detector
                checkpoint, as an alternative to the Hugging Face one

        Returns:
            The built WeightInitialization.
        """
        from deeplabcut.modelzoo import build_weight_init

        deprecation_warning = (
            "The `WeightInitialization.build` is deprecated and will be removed in a "
            "future version of DeepLabCut. Please use `build_weight_init` from "
            "`deeplabcut.modelzoo` instead."
        )
        warnings.warn(deprecation_warning, DeprecationWarning, stacklevel=2)

        return build_weight_init(
            cfg,
            super_animal,
            model_name,
            detector_name,
            with_decoder,
            memory_replay,
            customized_pose_checkpoint,
            customized_detector_checkpoint,
        )

build staticmethod 

build(
    cfg: dict,
    super_animal: str,
    model_name: str = "hrnet_w32",
    detector_name: str = "fasterrcnn_resnet50_fpn_v2",
    with_decoder: bool = False,
    memory_replay: bool = False,
    customized_pose_checkpoint: str | None = None,
    customized_detector_checkpoint: str | None = None,
) -> WeightInitialization

Builds a WeightInitialization for a project.

WeightInitialization.build is deprecated and will be removed in a future version of DeepLabCut. Please use build_weight_init from deeplabcut.modelzoo instead.

Parameters:

Name Type Description Default

cfg

dict

The project's configuration.

 required

super_animal

str

The SuperAnimal model with which to initialize weights.

 required

model_name

str

The name of the model architecture for which to load the weights (defaults to "hrnet_w32" for backwards compatibility).

 'hrnet_w32'

detector_name

str

The name of the detector architecture for which to load the weights (defaults to "fasterrcnn_resnet50_fpn_v2" for backwards compatibility).

 'fasterrcnn_resnet50_fpn_v2'

with_decoder

bool

Whether to load the decoder weights as well. If this is true, a conversion table must be specified for the given SuperAnimal in the project configuration file. See deeplabcut.modelzoo.utils.create_conversion_table to create a conversion table.

 False

memory_replay

bool

Only when with_decoder=True. Whether to train the model with memory replay, so that it predicts all SuperAnimal bodyparts.

 False

customized_pose_checkpoint

str | None

A customized SuperAnimal pose checkpoint, as an alternative to the Hugging Face one

 None

customized_detector_checkpoint

str | None

A customized SuperAnimal detector checkpoint, as an alternative to the Hugging Face one

 None

Returns:

Type Description
WeightInitialization

The built WeightInitialization.
Source code in deeplabcut/core/weight_init.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
@staticmethod
def build(
    cfg: dict,
    super_animal: str,
    model_name: str = "hrnet_w32",
    detector_name: str = "fasterrcnn_resnet50_fpn_v2",
    with_decoder: bool = False,
    memory_replay: bool = False,
    customized_pose_checkpoint: str | None = None,
    customized_detector_checkpoint: str | None = None,
) -> WeightInitialization:
    """Builds a WeightInitialization for a project.

    `WeightInitialization.build` is deprecated and will be removed in a future
    version of DeepLabCut. Please use `build_weight_init` from `deeplabcut.modelzoo`
    instead.

    Args:
        cfg: The project's configuration.
        super_animal: The SuperAnimal model with which to initialize weights.
        model_name: The name of the model architecture for which to load the weights
            (defaults to "hrnet_w32" for backwards compatibility).
        detector_name: The name of the detector architecture for which to load the
            weights (defaults to "fasterrcnn_resnet50_fpn_v2" for backwards
            compatibility).
        with_decoder: Whether to load the decoder weights as well. If this is true,
            a conversion table must be specified for the given SuperAnimal in the
            project configuration file. See
            ``deeplabcut.modelzoo.utils.create_conversion_table`` to create a
            conversion table.
        memory_replay: Only when ``with_decoder=True``. Whether to train the model
            with memory replay, so that it predicts all SuperAnimal bodyparts.
        customized_pose_checkpoint: A customized SuperAnimal pose checkpoint, as an
            alternative to the Hugging Face one
        customized_detector_checkpoint: A customized SuperAnimal detector
            checkpoint, as an alternative to the Hugging Face one

    Returns:
        The built WeightInitialization.
    """
    from deeplabcut.modelzoo import build_weight_init

    deprecation_warning = (
        "The `WeightInitialization.build` is deprecated and will be removed in a "
        "future version of DeepLabCut. Please use `build_weight_init` from "
        "`deeplabcut.modelzoo` instead."
    )
    warnings.warn(deprecation_warning, DeprecationWarning, stacklevel=2)

    return build_weight_init(
        cfg,
        super_animal,
        model_name,
        detector_name,
        with_decoder,
        memory_replay,
        customized_pose_checkpoint,
        customized_detector_checkpoint,
    )

from_dict_legacy classmethod 

from_dict_legacy(data: dict) -> Self

Deals with weight initialization that were created before 3.0.0rc5

Source code in deeplabcut/core/weight_init.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
@classmethod
def from_dict_legacy(cls, data: dict) -> Self:
    """Deals with weight initialization that were created before 3.0.0rc5"""

    import deeplabcut.pose_estimation_pytorch.modelzoo.utils as utils

    conversion_array = data.get("conversion_array")
    if conversion_array is not None:
        conversion_array = np.array(conversion_array, dtype=int)

    return cls(
        snapshot_path=utils.get_super_animal_snapshot_path(
            dataset=data["dataset"],
            model_name="hrnet_w32",
        ),
        detector_snapshot_path=utils.get_super_animal_snapshot_path(
            dataset=data["dataset"],
            model_name="fasterrcnn_resnet50_fpn_v2",
        ),
        with_decoder=data["with_decoder"],
        memory_replay=data["memory_replay"],
        conversion_array=conversion_array,
        bodyparts=data.get("bodyparts"),
    )