o  i@sddlmZddlZddlmZmZddlmZmZm Z ddl Z ddl m Z ddl mZddlmZddlmZmZmZmZdd lmZdd lmZmZdd lmZdd lmZmZm Z m!Z!m"Z"m#Z#dd l$m%Z&ddl$m'Z(erddl)m*Z*m+Z+ddl,m-Z-ne#de!j.e"d\Z*Z/e#de!j.e"d\Z-Z/e#de!j.e"d\Z+Z/gdZ0GdddeZ1Gddde1Z2Gddde1Z3Gddde1Z4dS)) annotationsN)IterableSequence) TYPE_CHECKINGAnyCallable) Optimizer) DataLoader) MetaTensor)IterationEventsdefault_make_latentdefault_metric_cmp_fndefault_prepare_batch)Workflow)Inferer SimpleInferer) Transform)AdversarialIterationEventsAdversarialKeysGanKeys IgniteInfo min_versionoptional_import) CommonKeys)EngineStatsKeys)Engine EventEnum)Metricz ignite.enginerzignite.metricsrr)TrainerSupervisedTrainer GanTrainerAdversarialTrainercs*eZdZdZdfdd ZddZZS) rzH Base class for all kinds of trainers, inherits from Workflow. returnNonecs&|jr tjjnd|_tdS)z Execute training based on Ignite Engine. If call this function multiple times, it will continuously run from the previous state. N)amptorchcuda GradScalerscalersuperrun)self __class__W/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/engines/trainer.pyr*0sz Trainer.runc GsVtj|jjtj|jjtj|jjtj|jj tj |jj i}|D] }t |j|d||<q|S)a Get the statistics information of the training process. Default to return the `rank`, `current_epoch`, `current_iteration`, `total_epochs`, `total_iterations`. Args: vars: except for the default stats, other variables name in the `self.state` to return, will use the variable name as the key and the state content as the value. if the variable doesn't exist, default value is `None`. N) ESKeysRANKstaterank CURRENT_EPOCHepochCURRENT_ITERATION iteration TOTAL_EPOCHS max_epochsTOTAL_ITERATIONS epoch_lengthgetattr)r+varsstatskr.r.r/ get_stats9s    zTrainer.get_statsr"r#)__name__ __module__ __qualname____doc__r*r@ __classcell__r.r.r,r/r*s rcsTeZdZdZddedddddeddddddddddfd7fd/d0 Zd8d5d6ZZS)9rac Standard supervised training method with image and label, inherits from ``Trainer`` and ``Workflow``. Args: device: an object representing the device on which to run. max_epochs: the total epoch number for trainer to run. train_data_loader: Ignite engine use data_loader to run, must be Iterable or torch.DataLoader. network: network to train in the trainer, should be regular PyTorch `torch.nn.Module`. optimizer: the optimizer associated to the network, should be regular PyTorch optimizer from `torch.optim` or its subclass. loss_function: the loss function associated to the optimizer, should be regular PyTorch loss, which inherit from `torch.nn.modules.loss`. epoch_length: number of iterations for one epoch, default to `len(train_data_loader)`. non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously with respect to the host. For other cases, this argument has no effect. prepare_batch: function to parse expected data (usually `image`, `label` and other network args) from `engine.state.batch` for every iteration, for more details please refer to: https://pytorch.org/ignite/generated/ignite.engine.create_supervised_trainer.html. iteration_update: the callable function for every iteration, expect to accept `engine` and `engine.state.batch` as inputs, return data will be stored in `engine.state.output`. if not provided, use `self._iteration()` instead. for more details please refer to: https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html. inferer: inference method that execute model forward on input data, like: SlidingWindow, etc. postprocessing: execute additional transformation for the model output data. Typically, several Tensor based transforms composed by `Compose`. key_train_metric: compute metric when every iteration completed, and save average value to engine.state.metrics when epoch completed. key_train_metric is the main metric to compare and save the checkpoint into files. additional_metrics: more Ignite metrics that also attach to Ignite Engine. metric_cmp_fn: function to compare current key metric with previous best key metric value, it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. train_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: CheckpointHandler, StatsHandler, etc. amp: whether to enable auto-mixed-precision training, default is False. event_names: additional custom ignite events that will register to the engine. new events can be a list of str or `ignite.engine.events.EventEnum`. event_to_attr: a dictionary to map an event to a state attribute, then add to `engine.state`. for more details, check: https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html #ignite.engine.engine.Engine.register_events. decollate: whether to decollate the batch-first data to a list of data after model computation, recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. default to `True`. optim_set_to_none: when calling `optimizer.zero_grad()`, instead of setting to zero, set the grads to None. more details: https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html. to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for `device`, `non_blocking`. amp_kwargs: dict of the args for `torch.autocast("cuda")` API, for more details: https://pytorch.org/docs/stable/amp.html#torch.autocast. compile: whether to use `torch.compile`, default is False. If True, MetaTensor inputs will be converted to `torch.Tensor` before forward pass, then converted back afterward with copied meta information. compile_kwargs: dict of the args for `torch.compile()` API, for more details: https://pytorch.org/docs/stable/generated/torch.compile.html#torch-compile. NFTdevicestr | torch.devicer9inttrain_data_loaderIterable | DataLoadernetworktorch.nn.Module optimizerr loss_functionrr; int | None non_blockingbool prepare_batchiteration_update#Callable[[Engine, Any], Any] | NoneinfererInferer | NonepostprocessingTransform | Nonekey_train_metricdict[str, Metric] | Noneadditional_metrics metric_cmp_fntrain_handlersSequence | Noner$ event_names.list[str | EventEnum | type[EventEnum]] | None event_to_attr dict | None decollateoptim_set_to_none to_kwargs amp_kwargscompilecompile_kwargsr"r#cstjdid|d|d|d|d|d| d| d| d | d |d |d |d |d|d|d|d|d||rP|durEin|}tj|fi|}||_||_||_||_| durctn| |_||_ dSNrGr9 data_loaderr;rQrSrTrX key_metricr\r]handlersr$r`rbrdrfrgr.) r)__init__r%rhrLrNrOrrVre)r+rGr9rJrLrNrOr;rQrSrTrVrXrZr\r]r^r$r`rbrdrerfrgrhrir,r.r/rns\       zSupervisedTrainer.__init__engine batchdatadict[str, torch.Tensor]dictc s0|durtdj|jjjfij}t|dkr&|\din|\|jr\d\}}}}tt rKt d j j}}tt r\ j j}}tjtjij_fdd}jjjjd jrjdurtjd ij |Wdn1swYjjjtj t!j"j#jj$n|jjtj t!j"j#|jr|durt ||d jjtj<t jjtj%||d jjtj%<|durt ||d jjtj< t!j&jjS) a Callback function for the Supervised Training processing logic of 1 iteration in Ignite Engine. Return below items in a dictionary: - IMAGE: image Tensor data for model input, already moved to device. - LABEL: label Tensor data corresponding to the image, already moved to device. - PRED: prediction result of model. - LOSS: loss value computed by loss function. Args: engine: `SupervisedTrainer` to execute operation for an iteration. batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. Raises: ValueError: When ``batchdata`` is None. N.Must provide batch data for current iteration.r.)NNNNzgWill convert to PyTorch Tensor if using compile, and casting back to MetaTensor after the forward pass.csfjjgRijjtj<tj jjtj jjtj <tj dSN) rVrLr2outputKeysPRED fire_eventr FORWARD_COMPLETEDrOmeanLOSSLOSS_COMPLETEDr.argsroinputskwargstargetsr.r/_compute_pred_losss& $z8SupervisedTrainer._iteration.._compute_pred_loss set_to_noner&)metaapplied_operationsr&)' ValueErrorrSr2rGrQrflenrh isinstancer warningswarn as_tensorrrrwIMAGELABELrvrLtrainrN zero_gradrer$r(r%autocastrgscaler|backwardryr BACKWARD_COMPLETEDstepupdaterxMODEL_COMPLETED) r+rorpbatch inputs_meta targets_metainputs_applied_operationstargets_applied_operationsrr.r~r/ _iterationsj             zSupervisedTrainer._iteration)4rGrHr9rIrJrKrLrMrNrrOrr;rPrQrRrSrrTrUrVrWrXrYrZr[r\r[r]rr^r_r$rRr`rarbrcrdrRrerRrfrcrgrcrhrRrircr"r#)rorrprqr"rr) rBrCrDrErr rnrrFr.r.r,r/rPs.?:rcsTeZdZdZddddddeedddddedddddfd:fd1d2 Zd;d8d9ZZ S)|jrs|jd||j |jj|jd|j}| ||j }|jj|jd||} | |jtj|tj|tj|tj| tj|iS)a Callback function for Adversarial Training processing logic of 1 iteration in Ignite Engine. Args: engine: `GanTrainer` to execute operation for an iteration. batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. Raises: ValueError: must provide batch data for current iteration. Nz.must provide batch data for current iteration.) num_latents latent_sizerGrQrrr.) rrSr2rGrQrfrk batch_sizerrrrr%zerosrangerrrrerrritemrrrrREALSFAKESLATENTSGLOSSDLOSS) r+rorpd_inputrZg_inputZg_outputZ d_total_loss_Zdlossg_lossr.r.r/rsP       zGanTrainer._iteration)8rGrHr9rIrJr rrMrrrrrrMrrrrr;rPrrWrrWrrIrrIrQrRrrrrrrRrTrUrXrYrZr[r\r[r]rr^r_rdrRrerRrfrcrgrc)ror rprr"r) rBrCrDrErr r rnrrFr.r.r,r/r s.HBr cs\eZdZdZddeddddddeddddddddfd<fd0d1 Zd=d4d5Zd>d:d;ZZ S)?r!a Standard supervised training workflow for adversarial loss enabled neural networks. Args: device: an object representing the device on which to run. max_epochs: the total epoch number for engine to run. train_data_loader: Core ignite engines uses `DataLoader` for training loop batchdata. g_network: ''generator'' (G) network architecture. g_optimizer: G optimizer function. g_loss_function: G loss function for adversarial training. recon_loss_function: G loss function for reconstructions. d_network: discriminator (D) network architecture. d_optimizer: D optimizer function. d_loss_function: D loss function for adversarial training.. epoch_length: number of iterations for one epoch, default to `len(train_data_loader)`. non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously with respect to the host. For other cases, this argument has no effect. prepare_batch: function to parse image and label for current iteration. iteration_update: the callable function for every iteration, expect to accept `engine` and `batchdata` as input parameters. if not provided, use `self._iteration()` instead. g_inferer: inference method to execute G model forward. Defaults to ``SimpleInferer()``. d_inferer: inference method to execute D model forward. Defaults to ``SimpleInferer()``. postprocessing: execute additional transformation for the model output data. Typically, several Tensor based transforms composed by `Compose`. Defaults to None key_train_metric: compute metric when every iteration completed, and save average value to engine.state.metrics when epoch completed. key_train_metric is the main metric to compare and save the checkpoint into files. additional_metrics: more Ignite metrics that also attach to Ignite Engine. metric_cmp_fn: function to compare current key metric with previous best key metric value, it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update 'best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. train_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: CheckpointHandler, StatsHandler, etc. amp: whether to enable auto-mixed-precision training, default is False. event_names: additional custom ignite events that will register to the engine. new events can be a list of str or `ignite.engine.events.EventEnum`. event_to_attr: a dictionary to map an event to a state attribute, then add to `engine.state`. for more details, check: https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html #ignite.engine.engine.Engine.register_events. decollate: whether to decollate the batch-first data to a list of data after model computation, recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. default to `True`. optim_set_to_none: when calling `optimizer.zero_grad()`, instead of setting to zero, set the grads to None. more details: https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html. to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for `device`, `non_blocking`. amp_kwargs: dict of the args for `torch.autocast("cuda")` API, for more details: https://pytorch.org/docs/stable/amp.html#torch.autocast. NFTrGtorch.device | strr9rIrJrKrrMrrrrrecon_loss_functionrrrr;rPrQrRrSrTCallable | NonerrWrrXrYrZr[r\r]r^r_r$r`rarbrcrdrerfrgcs*tjdid|d|d|d| d| d| d|d|d |d |d |d |d |d|d|d|d|d||jt||j_||j_||j_||j_||j_ | |j_ | |j_ |duret n||_ |durot n||_|jr{tjjnd|j_|jrtjjnd|j_||_|dSrj)r)rnregister_eventsrr2rrrrrrrrrrr$r%r&r'g_scalerd_scalerre_complete_state_dict_user_keys)r+rGr9rJrrrrrrrr;rQrSrTrrrXrZr\r]r^r$r`rbrdrerfrgr,r.r/rnsf        zAdversarialTrainer.__init__r"r#cCs|jgdt|jjdd}t|r|jdt|jjdd}t|r,|jdt|jjdd}t|r@|jddSdS)a1 This method appends to the _state_dict_user_keys AdversarialTrainer's elements that are required for checkpoint saving. Follows the example found at: https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html#ignite.engine.engine.Engine.state_dict )rrrrrr state_dictNrrr) Z_state_dict_user_keysextendr<r2rcallableappendrr)r+Zg_loss_state_dictZd_loss_state_dictZrecon_loss_state_dictr.r.r/rOs  z1AdversarialTrainer._complete_state_dict_user_keysrorprqrcs|durtdj|jjjfij}t|dkr&|\din|\tjtj t j ij_ d fdd }jj jjjjd jrjjdurtjdij |Wdn1sswYjj t jjj t jjj tj<jjjj tjtjjjjjjj n|jj t jjj t jtjjjtj!d fd d }jj"jj"jjd jr;jj#dur;tjdij |Wdn 1swYjj#jj t j$tj%jj#jj&jj# jj S|jj t j$jj&jj S)a Callback function for the Adversarial Training processing logic of 1 iteration in Ignite Engine. Return below items in a dictionary: - IMAGE: image Tensor data for model input, already moved to device. - LABEL: label Tensor data corresponding to the image, already moved to device. In case of Unsupervised Learning this is equal to IMAGE. - PRED: prediction result of model. - LOSS: loss value computed by loss functions of the generator (reconstruction and adversarial summed up). - AdversarialKeys.REALS: real images from the batch. Are the same as IMAGE. - AdversarialKeys.FAKES: fake images generated by the generator. Are the same as PRED. - AdversarialKeys.REAL_LOGITS: logits of the discriminator for the real images. - AdversarialKeys.FAKE_LOGITS: logits of the discriminator for the fake images. - AdversarialKeys.RECONSTRUCTION_LOSS: loss value computed by the reconstruction loss function. - AdversarialKeys.GENERATOR_LOSS: loss value computed by the generator loss function. It is the discriminator loss for the fake images. That is backpropagated through the generator only. - AdversarialKeys.DISCRIMINATOR_LOSS: loss value computed by the discriminator loss function. It is the discriminator loss for the real images and the fake images. That is backpropagated through the discriminator only. Args: engine: `AdversarialTrainer` to execute operation for an iteration. batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. Raises: ValueError: must provide batch data for current iteration. Nrsrtr.r"r#csjjjgRijjtj<jjtjjjtj<t j j jjtj jjgRijjtj<t jjjjtjjjtj<t jjjjtjjjtj<t jdSru)rr2rrvrrrwrxryrGENERATOR_FORWARD_COMPLETEDrfloat contiguousr FAKE_LOGITS)GENERATOR_DISCRIMINATOR_FORWARD_COMPLETEDrr{RECONSTRUCTION_LOSSRECONSTRUCTION_LOSS_COMPLETEDrGENERATOR_LOSSGENERATOR_LOSS_COMPLETEDr.r~r.r/_compute_generator_losss:     z>AdversarialTrainer._iteration.._compute_generator_lossrr&csjjjtjjjgRijjtj< t j jjjtj jjgRijjtj < t jjjjtjjjtj jjtj< t jdSru)rr2rvrrrdetachr REAL_LOGITSryr%DISCRIMINATOR_REALS_FORWARD_COMPLETEDrr%DISCRIMINATOR_FAKES_FORWARD_COMPLETEDrr{DISCRIMINATOR_LOSSDISCRIMINATOR_LOSS_COMPLETEDr.)rrorr.r/_compute_discriminator_losss0   zBAdversarialTrainer._iteration.._compute_discriminator_lossrAr)'rrSr2rGrQrfrrwrrrrrvrrrrrer$rr%rrgrrr|rrryrGENERATOR_BACKWARD_COMPLETEDrrGENERATOR_MODEL_COMPLETEDrrr DISCRIMINATOR_BACKWARD_COMPLETEDr)r+rorprrrr.r~r/rgs`                zAdversarialTrainer._iteration)8rGrr9rIrJrKrrMrrrrrrrrMrrrrr;rPrQrRrSrrTrrrWrrWrXrYrZr[r\r[r]rr^r_r$rRr`rarbrcrdrRrerRrfrcrgrcrA)ror!rprqr"r) rBrCrDrErr rnrrrFr.r.r,r/r!s.< Hr!)5 __future__rrcollections.abcrrtypingrrrr%Ztorch.optim.optimizerrtorch.utils.datar monai.datar monai.engines.utilsr r r rmonai.engines.workflowrmonai.inferersrrmonai.transformsr monai.utilsrrrrrrmonai.utils.enumsrrwrr0 ignite.enginerrignite.metricsrOPT_IMPORT_VERSIONr__all__rrr r!r.r.r.r/s8         &J>