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 dataclass

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

Parameters:

Name Type Description Default

snapshot_path

Path

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

 required

detector_snapshot_path

Path | None

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

 None

dataset

str | None

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

 None

with_decoder

bool

Whether to load the decoder weights as well.

 False

memory_replay

bool

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

 False

conversion_array

ndarray | None

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

 None

bodyparts

list[str] | None

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

 None

Methods:

Name Description
build

Builds a WeightInitialization for a project.
from_dict_legacy

Deals with weight initialization that were created before 3.0.0rc5.
to_dict

Returns: the weight initialization as a dict
Source code in deeplabcut/core/weight_init.py 
 22
 23
 24
 25
 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
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
@dataclass
class WeightInitialization:
    """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
    detector_snapshot_path: Path | None = None
    dataset: str | None = None
    with_decoder: bool = False
    memory_replay: bool = False
    conversion_array: np.ndarray | None = None
    bodyparts: list[str] | None = None

    def __post_init__(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 not len(self.conversion_array) == len(self.bodyparts):
                raise ValueError(
                    "There must be the same number of elements in the bodyparts list "
                    "and conv. array; found {self.bodyparts}, {self.conversion_array}"
                )

    def to_dict(self) -> dict:
        """Returns: the weight initialization as a dict"""
        data = dict()
        if self.dataset is not None:
            data["dataset"] = self.dataset

        data["snapshot_path"] = str(self.snapshot_path)
        if self.detector_snapshot_path is not None:
            data["detector_snapshot_path"] = str(self.detector_snapshot_path)

        data["with_decoder"] = self.with_decoder
        data["memory_replay"] = self.memory_replay

        if self.conversion_array is not None:
            data["conversion_array"] = self.conversion_array.tolist()

        if self.bodyparts is not None:
            data["bodyparts"] = self.bodyparts

        return data

    @staticmethod
    def from_dict(data: dict) -> WeightInitialization:
        if "snapshot_path" not in data:
            return WeightInitialization.from_dict_legacy(data)

        detector_snapshot_path = data.get("detector_snapshot_path")
        if detector_snapshot_path is not None:
            detector_snapshot_path = Path(detector_snapshot_path)

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

        return WeightInitialization(
            snapshot_path=Path(data["snapshot_path"]),
            detector_snapshot_path=detector_snapshot_path,
            dataset=data.get("dataset"),
            with_decoder=data["with_decoder"],
            memory_replay=data["memory_replay"],
            conversion_array=conversion_array,
            bodyparts=data.get("bodyparts"),
        )

    @staticmethod
    def from_dict_legacy(data: dict) -> WeightInitialization:
        """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 WeightInitialization(
            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 
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
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
@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 staticmethod 

from_dict_legacy(data: dict) -> WeightInitialization

Deals with weight initialization that were created before 3.0.0rc5.

Source code in deeplabcut/core/weight_init.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
@staticmethod
def from_dict_legacy(data: dict) -> WeightInitialization:
    """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 WeightInitialization(
        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"),
    )

to_dict 

to_dict() -> dict

Returns: the weight initialization as a dict

Source code in deeplabcut/core/weight_init.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def to_dict(self) -> dict:
    """Returns: the weight initialization as a dict"""
    data = dict()
    if self.dataset is not None:
        data["dataset"] = self.dataset

    data["snapshot_path"] = str(self.snapshot_path)
    if self.detector_snapshot_path is not None:
        data["detector_snapshot_path"] = str(self.detector_snapshot_path)

    data["with_decoder"] = self.with_decoder
    data["memory_replay"] = self.memory_replay

    if self.conversion_array is not None:
        data["conversion_array"] = self.conversion_array.tolist()

    if self.bodyparts is not None:
        data["bodyparts"] = self.bodyparts

    return data