o  ix@sddlmZddlZddlZddlZddlZddlmZddlm Z ddl m Z ddl Z ddl mZddlmZdgZGd ddZdS) ) annotationsN)Hashable) ModuleType)Any)DEFAULT_PROTOCOL)PathLike StateCacherc@sDeZdZdZddeefdddZ ddddZd ddZddZ 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")) NT in_memorybool cache_dirPathLike | Noneallow_overwrite pickle_modulerpickle_protocolintreturnNonecCsN||_|dur tn||_tj|jstd||_||_ ||_ i|_ dS)aQConstructor. 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: specifies pickle protocol when saving, with `torch.save`. Defaults to torch.serialization.DEFAULT_PROTOCOL. 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 ValueErrorr rrcached)selfr r r rrrZ/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/utils/state_cacher.py__init__-s zStateCacher.__init__keyrdata_objrModuleType | None int | NonecCs||jvr |js td|jr|j|dt|iidStj |j d|dt |d}|j|d|iit j |||durD|jn||durL|jn|dt|dr_|j|j|d<dSdS) 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)r"frrdevice)rr RuntimeErrorr updatecopydeepcopyrrjoinr idtorchsaverrhasattrr%)rrrrrfnrrrstoreRs" zStateCacher.storecCs||jvr td|d|jr|j|dS|j|d}tj|s,td|dtj|dddd }d |j|vrG| |j|d }|S) z2Retrieve the object stored under a given key name.zTarget z was not cached.r"zFailed to load state in z. File doesn't exist anymore.cSs|S)Nr)storagelocationrrrsz&StateCacher.retrieve..T) map_location weights_onlyr%) rKeyErrorr rrexistsr&r,loadto)rrr/rrrrretrievets  zStateCacher.retrievecCsD|js|jD]}tj|j|drt|j|dqdSdS)z>If necessary, delete any cached files existing in `cache_dir`.r"N)r rrrr7remove)rkrrr__del__s zStateCacher.__del__) r r r r r r rrrrrr)NN) rrrrrr rr!rr)rrrr) __name__ __module__ __qualname____doc__picklerrr0r:r=rrrrrs & " ) __future__rr(rrBrcollections.abcrtypesrtypingrr,Ztorch.serializationrmonai.config.type_definitionsr__all__rrrrrs