U PhC@sddlmZddlZddlZddlZddlZddlmZddlm Z m Z ddl Z ddl m Z ddlmZdgZGdddZdS) ) annotationsN) ModuleType)AnyHashable)DEFAULT_PROTOCOL)PathLike StateCacherc@sbeZdZdZddeefddddddd d d Zdd d ddddddZd d dddZddZ dS)raClass 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")) NTboolzPathLike | NonerintNone) in_memory cache_dirallow_overwrite pickle_modulepickle_protocolreturncCsN||_|dkrtn||_tj|js2td||_||_ ||_ i|_ dS)aJConstructor. 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. Nz+Given `cache_dir` is not a valid directory.) r tempfile gettempdirr ospathisdir ValueErrorrrrcached)selfr r rrrrM/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/utils/state_cacher.py__init__,szStateCacher.__init__rrzModuleType | Nonez int | None)keydata_objrrrcCs||jkr|jstd|jr:|j|dt|iin~tj |j d|dt |d}|j|d|iit j |||dkr|jn||dkr|jn|dt|dr|j|j|d<dS) a 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. z6Cached key already exists and overwriting is disabled.objZstate__z.ptN)rfrrdevice)rr RuntimeErrorr updatecopydeepcopyrrjoinr idtorchsaverrhasattrr")rrrrrfnrrrstoreQs" zStateCacher.store)rrcCs||jkrtd|d|jr.|j|dS|j|d}tj|sXtd|dtj|ddd}d |j|kr| |j|d }|S) z2Retrieve the object stored under a given key name.zTarget z was not cached.rzFailed to load state in z. File doesn't exist anymore.cSs|S)Nr)storagelocationrrr~z&StateCacher.retrieve..) map_locationr") rKeyErrorr rrexistsr#r)loadto)rrr,rrrrretrievess  zStateCacher.retrievecCs@|js<|jD].}tj|j|dr t|j|dq dS)z>If necessary, delete any cached files existing in `cache_dir`.rN)r rrrr4remove)rkrrr__del__s zStateCacher.__del__)NN) __name__ __module__ __qualname____doc__picklerrr-r7r:rrrrrs&") __future__rr%rr?rtypesrtypingrrr)Ztorch.serializationrZmonai.config.type_definitionsr__all__rrrrr s