refactored weight init for better modularity

n-poulsen · Oct 18, 2024 · Sep 18, 2024 · Oct 1, 2024 · Oct 1, 2024 · Oct 1, 2024
commit d6ee4167bf3d54b30a69c015e9cbf77a787002a3
diff --git a/deeplabcut/core/weight_init.py b/deeplabcut/core/weight_init.py
@@ -11,63 +11,46 @@
 """Classes to configure how to initialize model weights"""
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass
+from pathlib import Path
 
 import numpy as np
 
-import deeplabcut.modelzoo.utils as modelzoo_utils
-
 
 @dataclass
 class WeightInitialization:
-    """The dataset from which to initialize weights
-
-    To build a WeightInitialization instance for a project using the conversion table
-    specified in the project configuration file, use:
-
-        ```
-        from pathlib import Path
-        from deeplabcut.utils.auxiliaryfunctions import read_config
-
-        project_cfg = read_config("/path/to/my/project/config.yaml")
-        super_animal = "superanimal_quadruped"
-        weight_init = WeightInitialization.build(
-            cfg=project_cfg,
-            super_animal="superanimal_quadruped",
-            with_decoder=True,
-            memory_replay=True,
-        )
-        ```
+    """Configures weights initialization when transfer learning or fine-tuning models
 
     Args:
-        dataset: The dataset on which the model weights were trained. Must be one of the
-            SuperAnimal weights.
+        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 bodyparts.
-        conversion_array: The mapping from SuperAnimal to project bodyparts. Required
-            when `with_decoder=True`.
+            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.
-        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
     """
 
-    dataset: str
+    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
-    customized_pose_checkpoint: str | None = None
-    customized_detector_checkpoint: str | None = None
 
     def __post_init__(self):
-        # check that the dataset exists; raises a ValueError if it doesn't
-        _ = modelzoo_utils.get_super_animal_project_cfg(self.dataset)
         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 "
@@ -98,50 +81,97 @@ def __post_init__(self):
     def to_dict(self) -> dict:
         """Returns: the weight initialization as a dict"""
         data = {
-            "dataset": self.dataset,
+            "snapshot_path": str(self.snapshot_path),
             "with_decoder": self.with_decoder,
             "memory_replay": self.memory_replay,
         }
 
+        if self.dataset is not None:
+            data["dataset"] = self.dataset
+
+        if self.detector_snapshot_path is not None:
+            data["detector_snapshot_path"] = self.detector_snapshot_path
+
         if self.conversion_array is not None:
             data["conversion_array"] = self.conversion_array.tolist()
-        if self.customized_pose_checkpoint is not None:
-            data["customized_pose_checkpoint"] = self.customized_pose_checkpoint
-        if self.customized_detector_checkpoint is not None:
-            data["customized_detector_checkpoint"] = self.customized_detector_checkpoint
+
+        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(
-            dataset=data["dataset"],
+            snapshot_path=utils.get_super_animal_snapshot_path(
+                dataset=data["dataset"],
+                model_name="hrnetw32",
+            ),
+            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,
-            customized_pose_checkpoint=data.get("customized_pose_checkpoint"),
-            customized_detector_checkpoint=data.get("customized_detector_checkpoint"),
+            bodyparts=data.get("bodyparts"),
         )
 
     @staticmethod
     def build(
         cfg: dict,
         super_animal: str,
+        model_name: str = "hrnetw32",
+        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 "hrnetw32" 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
@@ -157,19 +187,21 @@ def build(
         Returns:
             The built WeightInitialization.
         """
-        conversion_array = None
-        bodyparts = None
-        if with_decoder:
-            conversion_table = modelzoo_utils.get_conversion_table(cfg, super_animal)
-            conversion_array = conversion_table.to_array()
-            bodyparts = conversion_table.converted_bodyparts()
-
-        return WeightInitialization(
-            dataset=super_animal,
-            with_decoder=with_decoder,
-            memory_replay=memory_replay,
-            conversion_array=conversion_array,
-            bodyparts=bodyparts,
-            customized_pose_checkpoint=customized_pose_checkpoint,
-            customized_detector_checkpoint=customized_detector_checkpoint,
+        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)
+
+        return build_weight_init(
+            cfg,
+            super_animal,
+            model_name,
+            detector_name,
+            with_decoder,
+            memory_replay,
+            customized_pose_checkpoint,
+            customized_detector_checkpoint,
         )
diff --git a/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py
@@ -587,7 +587,8 @@ def create_multianimaltraining_dataset(
                     pytorch_cfg = make_super_animal_finetune_config(
                         project_config=cfg,
                         pose_config_path=path_train_config,
-                        net_type=net_type,
+                        model_name=net_type,
+                        detector_name=detector_type,
                         weight_init=weight_init,
                     )
                 else:
@@ -603,7 +604,8 @@ def create_multianimaltraining_dataset(
                 auxiliaryfunctions.write_plainconfig(pose_cfg_path, pytorch_cfg)
 
             print(
-                "The training dataset is successfully created. Use the function 'train_network' to start training. Happy training!"
+                "The training dataset is successfully created. Use the function "
+                "'train_network' to start training. Happy training!"
             )
         else:
             pass

diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py
@@ -1275,7 +1275,8 @@ def create_training_dataset(
                         pytorch_cfg = make_super_animal_finetune_config(
                             project_config=cfg,
                             pose_config_path=path_train_config,
-                            net_type=net_type,
+                            model_name=net_type,
+                            detector_name=detector_type,
                             weight_init=weight_init,
                         )
                     else:

diff --git a/deeplabcut/gui/tabs/create_training_dataset.py b/deeplabcut/gui/tabs/create_training_dataset.py
@@ -32,6 +32,7 @@
     _create_label_widget,
 )
 from deeplabcut.gui.widgets import launch_napari
+from deeplabcut.modelzoo import build_weight_init
 from deeplabcut.utils.auxiliaryfunctions import (
     get_data_and_metadata_filenames,
     get_training_set_folder,
@@ -179,7 +180,6 @@ def log_augmentation_choice(self, augmentation):
 
     def edit_conversion_table(self):
         # Test beforehand whether a conversion table exists
-        weight_init = self.weight_init_selector.get_weight_init()
         memory_replay_folder = Path(self.root.project_folder) / "memory_replay"
         conversion_matrix_out_path = str(memory_replay_folder / "confusion_matrix.png")
         files = [self.root.config]
@@ -213,12 +213,6 @@ def create_training_dataset(self):
                 self.root.writer.write("Training dataset creation failed.")
                 return
 
-        try:
-            weight_init = self.weight_init_selector.get_weight_init()
-        except ValueError as err:
-            print(f"The training dataset could not be created: {err}.")
-            return
-
         if self.model_comparison:
             raise NotImplementedError
             # TODO: finish model_comparison
@@ -240,6 +234,14 @@ def create_training_dataset(self):
                 elif engine == Engine.PYTORCH and "top_down" in net_type:
                     detector_type = self.detector_choice.currentText()
 
+                try:
+                    weight_init = self.weight_init_selector.get_super_animal_weight_init(
+                        net_type, detector_type,
+                    )
+                except ValueError as err:
+                    print(f"The training dataset could not be created: {err}.")
+                    return
+
                 if self.data_split_selection.selected:
                     deeplabcut.create_training_dataset_from_existing_split(
                         self.root.config,
@@ -551,8 +553,16 @@ def update_choices(self, choices: list[str]) -> None:
             self.weight_init_choice.removeItem(0)
         self.weight_init_choice.addItems(choices)
 
-    def get_weight_init(self) -> WeightInitialization | None:
+    def get_super_animal_weight_init(
+        self, net_type: str, detector_type: str,
+    ) -> WeightInitialization | None:
         """
+        Args:
+            net_type: The architecture of the pose model from which to fine-tune a
+                SuperAnimal model.
+            detector_type: The architecture of the detector from which to fine-tune a
+                SuperAnimal model.
+
         Raises:
             ValueError if WeightInitialization should be defined but could not be
                 created (e.g. if there's no conversion table).
@@ -567,8 +577,10 @@ def get_weight_init(self) -> WeightInitialization | None:
         weight_init_data = _WEIGHT_INIT_OPTIONS[weight_init_choice]
         super_animal = weight_init_data["super_animal"]
         try:
-            weight_init = WeightInitialization.build(
+            weight_init = build_weight_init(
                 self.root.cfg,
+                model_name=net_type,
+                detector_name=detector_type,
                 super_animal=super_animal,
                 with_decoder=self.with_decoder,
                 memory_replay=self.memory_replay,

diff --git a/deeplabcut/modelzoo/__init__.py b/deeplabcut/modelzoo/__init__.py
@@ -8,3 +8,4 @@
 #
 # Licensed under GNU Lesser General Public License v3.0
 #
+from deeplabcut.modelzoo.weight_initialization import build_weight_init
diff --git a/deeplabcut/modelzoo/model_configs/fasterrcnn_resnet50_fpn_v2.yaml b/deeplabcut/modelzoo/model_configs/fasterrcnn_resnet50_fpn_v2.yaml
@@ -0,0 +1,35 @@
+data:
+  colormode: RGB
+  inference:
+    normalize_images: true
+  train:
+    hflip: true
+    normalize_images: true
+device: auto
+model:
+  type: FasterRCNN
+  variant: fasterrcnn_resnet50_fpn_v2
+  box_score_thresh: 0.6
+  pretrained: false
+runner:
+  type: DetectorTrainingRunner
+  eval_interval: 50
+  optimizer:
+    type: AdamW
+    params:
+      lr: 1e-5
+  scheduler:
+    type: LRListScheduler
+    params:
+      milestones: [ 90 ]
+      lr_list: [ [ 1e-6 ] ]
+  snapshots:
+    max_snapshots: 5
+    save_epochs: 50
+    save_optimizer_state: false
+train_settings:
+  batch_size: 1
+  dataloader_workers: 0
+  dataloader_pin_memory: true
+  display_iters: 500
+  epochs: 250
-Original file line number
+Diff line change
@@ Expand Up / @@ -8,3 +8,4 @@ @@
     #
     # Licensed under GNU Lesser General Public License v3.0
     #
+    from deeplabcut.modelzoo.weight_initialization import build_weight_init