Skip to content

Improved loading of snapshot weights with torch.load(..., weights_only=True) #2823

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MMathisLab merged 3 commits into from
Dec 23, 2024
Merged

Improved loading of snapshot weights with torch.load(..., weights_only=True) #2823

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
51ccedd
improved handling of weights_only
n-poulsen Dec 23, 2024
bc33cbf
fix edited file
n-poulsen Dec 23, 2024
5ce7e3f
fix error str
n-poulsen Dec 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions deeplabcut/pose_estimation_pytorch/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
8 changes: 4 additions & 4 deletions deeplabcut/pose_estimation_pytorch/apis/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
),
)

Expand All @@ -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

Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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}")
Expand Down
10 changes: 8 additions & 2 deletions deeplabcut/pose_estimation_pytorch/runners/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
137 changes: 119 additions & 18 deletions deeplabcut/pose_estimation_pytorch/runners/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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 `False`, 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()
26 changes: 15 additions & 11 deletions deeplabcut/pose_estimation_pytorch/runners/inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
26 changes: 15 additions & 11 deletions deeplabcut/pose_estimation_pytorch/runners/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__(
Expand All @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
1 change: 0 additions & 1 deletion examples/openfield-Pranav-2018-10-30/config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ identity:
# Project path (change when moving around)
project_path: WILL BE AUTOMATICALLY UPDATED BY DEMO CODE


# Default DeepLabCut engine to use for shuffle creation (either pytorch or tensorflow)
engine: pytorch

Expand Down
Loading