Class WeightInitialization

deeplabcut/core/weight_init.py:23–207

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

@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}"
                )