From 51ccedd2e0601711db8307bd77ef0fd9827879ea Mon Sep 17 00:00:00 2001 From: Niels Poulsen Date: Mon, 23 Dec 2024 20:52:40 +0100 Subject: [PATCH 1/3] improved handling of weights_only --- .../pose_estimation_pytorch/__init__.py | 4 + .../pose_estimation_pytorch/apis/utils.py | 8 +- .../runners/__init__.py | 10 +- .../pose_estimation_pytorch/runners/base.py | 137 +++++++++++++++--- .../runners/inference.py | 26 ++-- .../pose_estimation_pytorch/runners/train.py | 26 ++-- .../openfield-Pranav-2018-10-30/config.yaml | 6 +- .../runners/test_runners.py | 10 ++ .../runners/test_runners_inference.py | 5 +- 9 files changed, 184 insertions(+), 48 deletions(-) diff --git a/deeplabcut/pose_estimation_pytorch/__init__.py b/deeplabcut/pose_estimation_pytorch/__init__.py index 6adacc5898..56a16d2b8d 100644 --- a/deeplabcut/pose_estimation_pytorch/__init__.py +++ b/deeplabcut/pose_estimation_pytorch/__init__.py @@ -29,6 +29,10 @@ PoseDatasetParameters, ) from deeplabcut.pose_estimation_pytorch.data.dlcloader import DLCLoader +from deeplabcut.pose_estimation_pytorch.runners.base import ( + get_load_weights_only, + set_load_weights_only, +) from deeplabcut.pose_estimation_pytorch.runners.snapshots import TorchSnapshotManager from deeplabcut.pose_estimation_pytorch.task import Task from deeplabcut.pose_estimation_pytorch.utils import fix_seeds diff --git a/deeplabcut/pose_estimation_pytorch/apis/utils.py b/deeplabcut/pose_estimation_pytorch/apis/utils.py index b8df7677ad..726f2a3c1a 100644 --- a/deeplabcut/pose_estimation_pytorch/apis/utils.py +++ b/deeplabcut/pose_estimation_pytorch/apis/utils.py @@ -549,7 +549,7 @@ def get_inference_runners( max_individuals=max_individuals, ), load_weights_only=model_config["detector"]["runner"].get( - "load_weights_only", True, + "load_weights_only", None, ), ) @@ -562,7 +562,7 @@ def get_inference_runners( preprocessor=pose_preprocessor, postprocessor=pose_postprocessor, dynamic=dynamic, - load_weights_only=model_config["runner"].get("load_weights_only", True), + load_weights_only=model_config["runner"].get("load_weights_only", None), ) return pose_runner, detector_runner @@ -612,7 +612,7 @@ def get_detector_inference_runner( batch_size=batch_size, preprocessor=preprocessor, postprocessor=postprocessor, - load_weights_only=det_cfg["runner"].get("load_weights_only", True), + load_weights_only=det_cfg["runner"].get("load_weights_only", None), ) if not isinstance(runner, DetectorInferenceRunner): @@ -699,7 +699,7 @@ def get_pose_inference_runner( preprocessor=pose_preprocessor, postprocessor=pose_postprocessor, dynamic=dynamic, - load_weights_only=model_config["runner"].get("load_weights_only", True), + load_weights_only=model_config["runner"].get("load_weights_only", None), ) if not isinstance(runner, PoseInferenceRunner): raise RuntimeError(f"Failed to build PoseInferenceRunner for {model_config}") diff --git a/deeplabcut/pose_estimation_pytorch/runners/__init__.py b/deeplabcut/pose_estimation_pytorch/runners/__init__.py index f78cc2b22e..e5a5922df6 100644 --- a/deeplabcut/pose_estimation_pytorch/runners/__init__.py +++ b/deeplabcut/pose_estimation_pytorch/runners/__init__.py @@ -9,15 +9,21 @@ # Licensed under GNU Lesser General Public License v3.0 # -from deeplabcut.pose_estimation_pytorch.runners.base import Runner +from deeplabcut.pose_estimation_pytorch.runners.base import ( + attempt_snapshot_load, + get_load_weights_only, + fix_snapshot_metadata, + Runner, + set_load_weights_only, +) from deeplabcut.pose_estimation_pytorch.runners.dynamic_cropping import DynamicCropper -from deeplabcut.pose_estimation_pytorch.runners.logger import LOGGER from deeplabcut.pose_estimation_pytorch.runners.inference import ( build_inference_runner, DetectorInferenceRunner, InferenceRunner, PoseInferenceRunner, ) +from deeplabcut.pose_estimation_pytorch.runners.logger import LOGGER from deeplabcut.pose_estimation_pytorch.runners.snapshots import TorchSnapshotManager from deeplabcut.pose_estimation_pytorch.runners.train import ( build_training_runner, diff --git a/deeplabcut/pose_estimation_pytorch/runners/base.py b/deeplabcut/pose_estimation_pytorch/runners/base.py index fa99566088..c6a57c8e92 100644 --- a/deeplabcut/pose_estimation_pytorch/runners/base.py +++ b/deeplabcut/pose_estimation_pytorch/runners/base.py @@ -10,16 +10,44 @@ # from __future__ import annotations +import logging +import os import pickle from abc import ABC from pathlib import Path from typing import Generic, TypeVar +import numpy as np import torch import torch.nn as nn ModelType = TypeVar("ModelType", bound=nn.Module) +_load_weights_only: bool = ( + os.getenv("TORCH_LOAD_WEIGHTS_ONLY", "true").lower() in ("true", "1") +) + + +def get_load_weights_only() -> bool: + """Gets the default value to use when loading snapshots with `torch.load(...)`. + + Returns: + The default `weights_only` value when loading snapshots using `torch.load(...)`. + """ + global _load_weights_only + return _load_weights_only + + +def set_load_weights_only(value: bool) -> None: + """Sets the default value to use when loading snapshots with `torch.load(...)`. + + Args: + value: The default `weights_only` value to use when loading snapshots using + `torch.load(...)`. + """ + global _load_weights_only + _load_weights_only = value + class Runner(ABC, Generic[ModelType]): """Runner base class @@ -63,7 +91,7 @@ def load_snapshot( snapshot_path: str | Path, device: str, model: ModelType, - weights_only: bool = True, + weights_only: bool | None = None, ) -> dict: """Loads the state dict for a model from a file @@ -74,11 +102,13 @@ def load_snapshot( snapshot_path: The path containing the model weights to load device: The device on which the model should be loaded model: The model for which the weights are loaded - weights_only: Value for torch.load() `weights_only` parameter. If False, the - python pickle module is used implicitly, which is known to be insecure. - Only set to False if you're loading data that you trust (e.g. snapshots - that you created yourself). For more information, see: + weights_only: Value for torch.load() `weights_only` parameter. + If False, the python pickle module is used implicitly, which is known to + be insecure. Only set to False if you're loading data that you trust + (e.g. snapshots that you created yourself). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html + If None, the default value is used: + `deeplabcut.pose_estimation_pytorch.get_load_weights_only()` Returns: The content of the snapshot file. @@ -88,17 +118,23 @@ def load_snapshot( return snapshot -def attempt_snapshot_load(path: str | Path, device: str, weights_only: bool) -> dict: +def attempt_snapshot_load( + path: str | Path, + device: str, + weights_only: bool | None = None, +) -> dict: """Attempts to load a snapshot using `torch.load(...)`. Args: path: The path of the snapshot to try to load.. device: The device to use for the `map_location`. - weights_only: Value for torch.load() `weights_only` parameter. If False, the - python pickle module is used implicitly, which is known to be insecure. - Only set to False if you're loading data that you trust (e.g. snapshots - that you created yourself). For more information, see: + weights_only: Value for torch.load() `weights_only` parameter. + If False, the python pickle module is used implicitly, which is known to be + insecure. Only set to False if you're loading data that you trust (e.g. + snapshots that you created yourself). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html + If None, the default value is used: + `deeplabcut.pose_estimation_pytorch.get_load_weights_only()` Returns: The loaded snapshot. @@ -108,18 +144,83 @@ def attempt_snapshot_load(path: str | Path, device: str, weights_only: bool) -> with `weights_only=True`. """ try: + if weights_only is None: + weights_only = get_load_weights_only() + snapshot = torch.load(path, map_location=device, weights_only=weights_only) except pickle.UnpicklingError as err: - print( - f"\nFailed to load the snapshot: {path}.\n" - "If you trust the snapshot that you're trying to load, you can try " - "calling `Runner.load_snapshot` with `weights_only=False`. See " - "the message below for more information and warnings.\n" - "You can set the `weights_only` parameter in the model configuration (" - "the content of the pytorch_config.yaml), as:\n```\n" + logging.error( + f"\nFailed to load the snapshot: {path}.\n\n" + "If you trust the snapshot that you're trying to load, you can try\n" + "calling `Runner.load_snapshot` with `weights_only=False`. See the \n" + "error message below for more information and warnings.\n" + "You can set the `weights_only` parameter in the model configuration (\n" + "the content of the pytorch_config.yaml), as:\n\n```\n" "runner:\n" - " load_weights_only: False\n```\n" + " load_weights_only: False\n```\n\n" + "If it's the detector snapshot that's failing to load, place the\n" + "`load_weights_only` key under the detector runner:\n\n```\n" + "detector:\n" + " runner:\n" + " load_weights_only: False\n```\n\n" + "You can also set the default `load_weights_only` that will be used when\n" + "the `load_weights_only` variable is not set in the `pytorch_config.yaml`\n" + "using `deeplabcut.pose_estimation_pytorch.set_load_weights_only(value)`:\n" + "\n```\n" + "from deeplabcut.pose_estimation_pytorch import set_load_weights_only\n" + "set_load_weights_only(True)\n" + "```\n\n" + "You can also set the value for `load_weights_only` with a \n" + "`TORCH_LOAD_WEIGHTS_ONLY` environment variable. If you call \n" + "`TORCH_LOAD_WEIGHTS_ONLY=False python -m deeplabcut`, it will launch the\n" + "DeepLabCut GUI with the default `load_weights_only` value to False.\n" + "If you set this value to `True`, make sure you only load snapshots that\n" + "you trust.\n\n" ) raise err return snapshot + + +def fix_snapshot_metadata(path: str | Path) -> None: + """Replace numpy floats in snapshot metrics + + Only call this method with snapshots that you trust, as torch.load(...) is called + with `weights_only=False`. For more information, see: + https://pytorch.org/docs/stable/generated/torch.load.html + + DeepLabCut PyTorch snapshots trained with older releases may have `numpy` floats in + the stored metrics. This method opens the snapshots (with `weights_only=False`), + replaces the numpy floats with python floats (allowing to load with + `weights_only=True`), and saves the new snapshot data. + + Warning: This overwrites your existing snapshot. If you want to ensure that no data + is lost, copy your snapshot before calling `fix_snapshot_metadata`. + + Args: + path: The path of the snapshot to fix. + """ + snapshot = torch.load(path, map_location="cpu", weights_only=False) + metrics = snapshot.get("metadata", {}).get("metrics") + if metrics is not None: + snapshot["metadata"]["metrics"] = {k: float(v) for k, v in metrics.items()} + + torch.save(snapshot, path) + + +def _add_numpy_to_torch_safe_globals(): + """ + Attempts tot add numpy classes allowing snapshots containing numpy floats in the + metrics to be loaded without needing to change the `weights_only` argument. + + This fix only works for `numpy>=1.25.0`. + """ + try: + from numpy.core.multiarray import scalar + from numpy.dtypes import Float64DType + torch.serialization.add_safe_globals([np.dtype, Float64DType, scalar]) + except Exception: + pass + + +_add_numpy_to_torch_safe_globals() diff --git a/deeplabcut/pose_estimation_pytorch/runners/inference.py b/deeplabcut/pose_estimation_pytorch/runners/inference.py index ecbc0a8a40..a013b0c780 100644 --- a/deeplabcut/pose_estimation_pytorch/runners/inference.py +++ b/deeplabcut/pose_estimation_pytorch/runners/inference.py @@ -42,7 +42,7 @@ def __init__( snapshot_path: str | Path | None = None, preprocessor: Preprocessor | None = None, postprocessor: Postprocessor | None = None, - load_weights_only: bool = True, + load_weights_only: bool | None = None, ): """ Args: @@ -52,11 +52,13 @@ def __init__( pretrained weights preprocessor: The preprocessor to use on images before inference postprocessor: The postprocessor to use on images after inference - load_weights_only: Value for the torch.load() `weights_only` parameter. If - False, the python pickle module is used implicitly, which is known to be - insecure. Only set to False if you're loading data that you trust (e.g. - snapshots that you created yourself). For more information, see: - https://pytorch.org/docs/stable/generated/torch.load.html + load_weights_only: Value for the torch.load() `weights_only` parameter. + If False, the python pickle module is used implicitly, which is known to + be insecure. Only set to False if you're loading data that you trust + (e.g. snapshots that you created). For more information, see: + https://pytorch.org/docs/stable/generated/torch.load.html + If None, the default value is used: + `deeplabcut.pose_estimation_pytorch.get_load_weights_only()` """ super().__init__(model=model, device=device, snapshot_path=snapshot_path) if not isinstance(batch_size, int) or batch_size <= 0: @@ -328,7 +330,7 @@ def build_inference_runner( preprocessor: Preprocessor | None = None, postprocessor: Postprocessor | None = None, dynamic: DynamicCropper | None = None, - load_weights_only: bool = True, + load_weights_only: bool | None = None, ) -> InferenceRunner: """ Build a runner object according to a pytorch configuration file @@ -345,11 +347,13 @@ def build_inference_runner( cropping should not be used. Only for bottom-up pose estimation models. Should only be used when creating inference runners for video pose estimation with batch size 1. - load_weights_only: Value for the torch.load() `weights_only` parameter. If - False, the python pickle module is used implicitly, which is known to be - insecure. Only set to False if you're loading data that you trust (e.g. - snapshots that you created yourself). For more information, see: + load_weights_only: Value for the torch.load() `weights_only` parameter. + If False, the python pickle module is used implicitly, which is known to + be insecure. Only set to False if you're loading data that you trust (e.g. + snapshots that you created). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html + If None, the default value is used: + `deeplabcut.pose_estimation_pytorch.get_load_weights_only()` Returns: The inference runner. diff --git a/deeplabcut/pose_estimation_pytorch/runners/train.py b/deeplabcut/pose_estimation_pytorch/runners/train.py index ed06e02abb..2be3e8ec17 100644 --- a/deeplabcut/pose_estimation_pytorch/runners/train.py +++ b/deeplabcut/pose_estimation_pytorch/runners/train.py @@ -68,11 +68,13 @@ class TrainingRunner(Runner, Generic[ModelType], metaclass=ABCMeta): logger: Logger to monitor training (e.g. a WandBLogger). log_filename: Name of the file in which to store training stats. load_weights_only: Value for the torch.load() `weights_only` parameter if - `snapshot_path` is not None. If False, the python pickle module is used - implicitly, which is known to be insecure. Only set to False if you're - loading data that you trust (e.g. snapshots that you created yourself). For - more information, see: + `snapshot_path` is not None. + If False, the python pickle module is used implicitly, which is known to + be insecure. Only set to False if you're loading data that you trust + (e.g. snapshots that you created yourself). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html + If None, the default value is used: + `deeplabcut.pose_estimation_pytorch.get_load_weights_only()` """ def __init__( @@ -88,7 +90,7 @@ def __init__( load_scheduler_state_dict: bool = True, logger: BaseLogger | None = None, log_filename: str = "learning_stats.csv", - load_weights_only: bool = True, + load_weights_only: bool | None = None, ): super().__init__( model=model, device=device, gpus=gpus, snapshot_path=snapshot_path @@ -368,7 +370,7 @@ def load_snapshot( snapshot_path: str | Path, device: str, model: PoseModel, - weights_only: bool = True, + weights_only: bool | None = None, ) -> dict: """Loads the state dict for a model from a file @@ -379,11 +381,13 @@ def load_snapshot( snapshot_path: the path containing the model weights to load device: the device on which the model should be loaded model: the model for which the weights are loaded - weights_only: Value for torch.load() `weights_only` parameter. If False, the - python pickle module is used implicitly, which is known to be insecure. - Only set to False if you're loading data that you trust (e.g. snapshots - that you created yourself). For more information, see: + weights_only: Value for torch.load() `weights_only` parameter. + If False, the python pickle module is used implicitly, which is known to + be insecure. Only set to False if you're loading data that you trust + (e.g. snapshots that you created yourself). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html + If None, the default value is used: + `deeplabcut.pose_estimation_pytorch.get_load_weights_only()` Returns: The content of the snapshot file. @@ -733,7 +737,7 @@ def build_training_runner( scheduler=scheduler, load_scheduler_state_dict=runner_config.get("load_scheduler_state_dict", True), logger=logger, - load_weights_only=runner_config.get("load_weights_only", True), + load_weights_only=runner_config.get("load_weights_only", None), ) if task == Task.DETECT: return DetectorTrainingRunner(**kwargs) diff --git a/examples/openfield-Pranav-2018-10-30/config.yaml b/examples/openfield-Pranav-2018-10-30/config.yaml index d6bdced20b..803b537e3d 100644 --- a/examples/openfield-Pranav-2018-10-30/config.yaml +++ b/examples/openfield-Pranav-2018-10-30/config.yaml @@ -7,7 +7,8 @@ identity: # Project path (change when moving around) -project_path: WILL BE AUTOMATICALLY UPDATED BY DEMO CODE +project_path: + /Users/niels/Documents/upamathis/repos/DeepLabCut/examples/openfield-Pranav-2018-10-30 # Default DeepLabCut engine to use for shuffle creation (either pytorch or tensorflow) @@ -25,6 +26,9 @@ bodyparts: - tailbase +# Fraction of video to start/stop when extracting frames for labeling/refinement + + # Fraction of video to start/stop when extracting frames for labeling/refinement start: 0 stop: 1 diff --git a/tests/pose_estimation_pytorch/runners/test_runners.py b/tests/pose_estimation_pytorch/runners/test_runners.py index 5d777ca410..3f2fc2e3da 100644 --- a/tests/pose_estimation_pytorch/runners/test_runners.py +++ b/tests/pose_estimation_pytorch/runners/test_runners.py @@ -19,11 +19,21 @@ import deeplabcut.pose_estimation_pytorch.runners as runners +@pytest.mark.parametrize("value", [True, False]) +def test_set_load_weights_only(value: bool): + print(f"\nget_load_weights_only: {runners.get_load_weights_only()}") + print(f"setting value to {value}") + runners.set_load_weights_only(value) + print(f"get_load_weights_only: {runners.get_load_weights_only()}\n") + assert runners.get_load_weights_only() == value + + def test_load_snapshot_weights_only_error(tmpdir_factory): snapshot_dir = Path(tmpdir_factory.mktemp("snapshot-dir")) snapshot_path = snapshot_dir / "snapshot.pt" torch.save(dict(content=np.zeros(10)), str(snapshot_path)) + runners.set_load_weights_only(False) with pytest.raises(pickle.UnpicklingError): runners.Runner.load_snapshot( snapshot_path, device="cpu", model=Mock(), weights_only=True diff --git a/tests/pose_estimation_pytorch/runners/test_runners_inference.py b/tests/pose_estimation_pytorch/runners/test_runners_inference.py index 3ca13a2e23..272b5eb509 100644 --- a/tests/pose_estimation_pytorch/runners/test_runners_inference.py +++ b/tests/pose_estimation_pytorch/runners/test_runners_inference.py @@ -18,12 +18,13 @@ import deeplabcut.pose_estimation_pytorch.data.postprocessor as post import deeplabcut.pose_estimation_pytorch.data.preprocessor as prep import deeplabcut.pose_estimation_pytorch.runners.inference as inference +from deeplabcut.pose_estimation_pytorch import get_load_weights_only from deeplabcut.pose_estimation_pytorch.task import Task @patch("deeplabcut.pose_estimation_pytorch.runners.train.build_optimizer", Mock()) @pytest.mark.parametrize("task", [Task.DETECT, Task.TOP_DOWN, Task.BOTTOM_UP]) -@pytest.mark.parametrize("weights_only", [True, False]) +@pytest.mark.parametrize("weights_only", [None, True, False]) def test_load_weights_only_with_build_training_runner(task: Task, weights_only: bool): with patch("deeplabcut.pose_estimation_pytorch.runners.base.torch.load") as load: snapshot = "snapshot.pt" @@ -34,6 +35,8 @@ def test_load_weights_only_with_build_training_runner(task: Task, weights_only: snapshot_path=snapshot, load_weights_only=weights_only, ) + if weights_only is None: + weights_only = get_load_weights_only() load.assert_called_once_with( snapshot, map_location="cpu", weights_only=weights_only ) From bc33cbf5eb0615e6e086f40cf7ff3f1b2b013241 Mon Sep 17 00:00:00 2001 From: Niels Poulsen Date: Mon, 23 Dec 2024 21:09:36 +0100 Subject: [PATCH 2/3] fix edited file --- examples/openfield-Pranav-2018-10-30/config.yaml | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/examples/openfield-Pranav-2018-10-30/config.yaml b/examples/openfield-Pranav-2018-10-30/config.yaml index 803b537e3d..ed8c31fdf3 100644 --- a/examples/openfield-Pranav-2018-10-30/config.yaml +++ b/examples/openfield-Pranav-2018-10-30/config.yaml @@ -7,9 +7,7 @@ identity: # Project path (change when moving around) -project_path: - /Users/niels/Documents/upamathis/repos/DeepLabCut/examples/openfield-Pranav-2018-10-30 - +project_path: WILL BE AUTOMATICALLY UPDATED BY DEMO CODE # Default DeepLabCut engine to use for shuffle creation (either pytorch or tensorflow) engine: pytorch @@ -26,9 +24,6 @@ bodyparts: - tailbase -# Fraction of video to start/stop when extracting frames for labeling/refinement - - # Fraction of video to start/stop when extracting frames for labeling/refinement start: 0 stop: 1 From 5ce7e3f7d8bb477571fab05a3f810a46c364e896 Mon Sep 17 00:00:00 2001 From: Niels Poulsen Date: Mon, 23 Dec 2024 21:40:12 +0100 Subject: [PATCH 3/3] fix error str --- deeplabcut/pose_estimation_pytorch/runners/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deeplabcut/pose_estimation_pytorch/runners/base.py b/deeplabcut/pose_estimation_pytorch/runners/base.py index c6a57c8e92..ee8dbd6d1a 100644 --- a/deeplabcut/pose_estimation_pytorch/runners/base.py +++ b/deeplabcut/pose_estimation_pytorch/runners/base.py @@ -174,7 +174,7 @@ def attempt_snapshot_load( "`TORCH_LOAD_WEIGHTS_ONLY` environment variable. If you call \n" "`TORCH_LOAD_WEIGHTS_ONLY=False python -m deeplabcut`, it will launch the\n" "DeepLabCut GUI with the default `load_weights_only` value to False.\n" - "If you set this value to `True`, make sure you only load snapshots that\n" + "If you set this value to `False`, make sure you only load snapshots that\n" "you trust.\n\n" ) raise err