o  i@sddlmZddlZddlZddlmZddlZddlmZddl m Z m Z m Z e de j e d\ZZe de j e d \ZZerEdd lmZn e de j e d \ZZGd d d ZdS)) annotationsN) TYPE_CHECKING)copy_model_state) IgniteInfo min_versionoptional_importz ignite.engineEventszignite.handlers Checkpoint)Enginer c@s8eZdZdZ    ddddZdddZdddZdS)CheckpointLoaderu CheckpointLoader acts as an Ignite handler to load checkpoint data from file. It can load variables for network, optimizer, lr_scheduler, etc. If saving checkpoint after `torch.nn.DataParallel`, need to save `model.module` instead as PyTorch recommended and then use this loader to load the model. Usage example:: trainer = SupervisedTrainer(...) save_dict = { "trainer": trainer, "net": network, "opt": optimizer, "lr": lr_scheduler, } map_location = "cuda:0" # checkpoint needs to have same save_dict for this to work handler = CheckpointLoader(load_path="/test/checkpoint.pt", load_dict=save_dict, map_location=map_location, strict=True) handler(trainer) # Trainer now has the same state as stored, including the number of epochs and iterations completed # so you can resume an interrupted training at the place where it left Args: load_path: the file path of checkpoint, it should be a PyTorch `pth` file. load_dict: target objects that load checkpoint to. examples:: {'network': net, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler} name: identifier of logging.logger to use, if None, defaulting to ``engine.logger``. map_location: when loading the module for distributed training/evaluation, need to provide an appropriate map_location argument to prevent a process to step into others’ devices. If map_location is missing, torch.load will first load the module to CPU and then copy each parameter to where it was saved, which would result in all processes on the same machine using the same set of devices. strict: whether to strictly enforce that the keys and data shape in the `state_dict` of every item of `load_dict` match the `state_dict` of the corresponding items of checkpoint, default to `True`. strict_shape: whether to enforce the data shape of the matched layers in the checkpoint, `if `False`, it will skip the layers that have different data shape with checkpoint content, and ignore the `strict` arg. this can be useful advanced feature for transfer learning. users should totally understand which layers will have different shape. default to `True`. Note: if `strict_shape=False`, will only load checkpoint for `torch.nn.Module` and skip other items in the `load_dict`. For example, if the shape of some layers in current model can't match the checkpoint, the `parameter_group` of current optimizer may also can't match the checkpoint, so skip loading checkpoint for optimizer. For more details about loading checkpoint, please refer to: https://pytorch.org/ignite/v0.4.5/generated/ignite.handlers.checkpoint.Checkpoint.html #ignite.handlers.checkpoint.Checkpoint.load_objects. https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.load_state_dict. NT load_pathstr load_dictdictname str | None map_location dict | Nonestrictbool strict_shapereturnNonecCsv|durtd||_|dust|dkrtdt||_||_||_||_|r3|s3t dd}||_ ||_ dS)Nz+must provide clear path to load checkpoint.rz$must provide target objects to load.z=as `strict_shape` is already False, change `strict` to False.F) AssertionErrorr lenlogging getLoggerloggerr_namerwarningswarnrr)selfr rrrrrr"b/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/handlers/checkpoint_loader.py__init__Ws    zCheckpointLoader.__init__enginer cCs$|jdur |j|_|tj|dS)g Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. N)rradd_event_handlerrSTARTED)r!r%r"r"r#attachos zCheckpointLoader.attachc Cs tj|j|jdd}t|jd\}}t|jdkr$||vr$||i}|js[g}|jD]!\}}t |tj j rEt ||ddd||<q.t d||q.|D]}|j|qR|jj}tj|j||jd|d ur|jj|krtd |jjd |d ||j_|jd |jd S)r&T)r weights_onlyrF)inplacezO`strict_shape` is False, load checkpoint for model, skip others in `load_dict`.)to_load checkpointrNz Epoch count (z>) in checkpoint is larger than the `engine.state.max_epochs` (z) of engine. To further train from checkpoint, construct trainer with `max_epochs` larger than checkpoint's epoch count. To use checkpoint for inference, no need to load state_dict for the engine.zRestored all variables from )torchloadr rlistritemsrr isinstancennModulerrr appendpopstate max_epochsr load_objectsrepoch ValueErrorrinfo) r!r%r.k_Z pop_itemsobjiZprior_max_epochsr"r"r#__call__xs.   zCheckpointLoader.__call__)NNTT)r r rrrrrrrrrrrr)r%r rr)__name__ __module__ __qualname____doc__r$r)rBr"r"r"r#r s;  r ) __future__rrrtypingrr/monai.networks.utilsr monai.utilsrrrOPT_IMPORT_VERSIONrr?r ignite.enginer r r"r"r"r#s