Source code for monai.utils.state_cacher
# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import os
import pickle
import tempfile
from typing import Dict, Optional
import torch
from torch.serialization import DEFAULT_PROTOCOL
from monai.config.type_definitions import PathLike
__all__ = ["StateCacher"]
[docs]class StateCacher:
"""Class to cache and retrieve the state of an object.
Objects can either be stored in memory or on disk. If stored on disk, they can be
stored in a given directory, or alternatively a temporary location will be used.
If necessary/possible, restored objects will be returned to their original device.
Example:
>>> state_cacher = StateCacher(memory_cache, cache_dir=cache_dir)
>>> state_cacher.store("model", model.state_dict())
>>> model.load_state_dict(state_cacher.retrieve("model"))
"""
[docs] def __init__(
self,
in_memory: bool,
cache_dir: Optional[PathLike] = None,
allow_overwrite: bool = True,
pickle_module=pickle,
pickle_protocol: int = DEFAULT_PROTOCOL,
) -> None:
"""Constructor.
Args:
in_memory: boolean to determine if the object will be cached in memory or on
disk.
cache_dir: directory for data to be cached if `in_memory==False`. Defaults
to using a temporary directory. Any created files will be deleted during
the `StateCacher`'s destructor.
allow_overwrite: allow the cache to be overwritten. If set to `False`, an
error will be thrown if a matching already exists in the list of cached
objects.
pickle_module: module used for pickling metadata and objects, default to `pickle`.
this arg is used by `torch.save`, for more details, please check:
https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
pickle_protocol: can be specified to override the default protocol, default to `2`.
this arg is used by `torch.save`, for more details, please check:
https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
"""
self.in_memory = in_memory
self.cache_dir = tempfile.gettempdir() if cache_dir is None else cache_dir
if not os.path.isdir(self.cache_dir):
raise ValueError("Given `cache_dir` is not a valid directory.")
self.allow_overwrite = allow_overwrite
self.pickle_module = pickle_module
self.pickle_protocol = pickle_protocol
self.cached: Dict = {}
[docs] def store(self, key, data_obj, pickle_module=None, pickle_protocol: Optional[int] = None):
"""
Store a given object with the given key name.
Args:
key: key of the data object to store.
data_obj: data object to store.
pickle_module: module used for pickling metadata and objects, default to `self.pickle_module`.
this arg is used by `torch.save`, for more details, please check:
https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
pickle_protocol: can be specified to override the default protocol, default to `self.pickle_protocol`.
this arg is used by `torch.save`, for more details, please check:
https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
"""
if key in self.cached and not self.allow_overwrite:
raise RuntimeError("Cached key already exists and overwriting is disabled.")
if self.in_memory:
self.cached.update({key: {"obj": copy.deepcopy(data_obj)}})
else:
fn = os.path.join(self.cache_dir, f"state_{key}_{id(self)}.pt")
self.cached.update({key: {"obj": fn}})
torch.save(
obj=data_obj,
f=fn,
pickle_module=self.pickle_module if pickle_module is None else pickle_module,
pickle_protocol=self.pickle_protocol if pickle_protocol is None else pickle_protocol,
)
# store object's device if relevant
if hasattr(data_obj, "device"):
self.cached[key]["device"] = data_obj.device
[docs] def retrieve(self, key):
"""Retrieve the object stored under a given key name."""
if key not in self.cached:
raise KeyError(f"Target {key} was not cached.")
if self.in_memory:
return self.cached[key]["obj"]
fn = self.cached[key]["obj"] # pytype: disable=attribute-error
if not os.path.exists(fn): # pytype: disable=wrong-arg-types
raise RuntimeError(f"Failed to load state in {fn}. File doesn't exist anymore.")
data_obj = torch.load(fn, map_location=lambda storage, location: storage)
# copy back to device if necessary
if "device" in self.cached[key]:
data_obj = data_obj.to(self.cached[key]["device"])
return data_obj
def __del__(self):
"""If necessary, delete any cached files existing in `cache_dir`."""
if not self.in_memory:
for k in self.cached:
if os.path.exists(self.cached[k]["obj"]):
os.remove(self.cached[k]["obj"])