U vPh\@s~ddlmZddlZddlmZmZmZmZmZddl Z ddl m Z ddl m Z ddlmZddlmZddlmZmZmZmZdd lmZdd lmZmZdd lmZdd lmZm Z m!Z!dd l"m#Z$ddl"m%Z&ddl'm(Z(erddl)m*Z*m+Z+ddl,m-Z-n __classcell__r,r,r*r-r+s cseZdZdZddedddddeddddddddddfddddd d d d d d ddddd dd ddd d ddd dddfdd ZddddddZZS)rao 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.cuda.amp.autocast()` API, for more details: https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.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. NFTstr | torch.deviceintzIterable | DataLoadertorch.nn.Modulerr int | Nonebool#Callable[[Engine, Any], Any] | NoneInferer | NoneTransform | Nonedict[str, Metric] | NoneSequence | Nonez.list[str | EventEnum | type[EventEnum]] | None dict | Noner )devicer7train_data_loadernetwork optimizer loss_functionr9 non_blocking prepare_batchiteration_updateinfererpostprocessingkey_train_metricadditional_metrics metric_cmp_fntrain_handlersr" event_names event_to_attr decollateoptim_set_to_none to_kwargs amp_kwargscompilecompile_kwargsr!cstj|||||| | | | |||||||||d|rhtddr^|dkrJin|}tj|f|}n td||_||_||_||_ | dkrt n| |_ ||_ dS)N)rOr7 data_loaderr9rTrUrVrX key_metricrZr[handlersr"r]r^r_rarbzeNetwork compilation (compile=True) not supported for Pytorch versions before 2.1, no compilation done) r'__init__rr#rcwarningswarnrQrRrSrrWr`)r)rOr7rPrQrRrSr9rTrUrVrWrXrYrZr[r\r"r]r^r_r`rarbrcrdr*r,r-rjsB zSupervisedTrainer.__init__zdict[str, torch.Tensor]dictengine batchdatar!c s|dkrtdj|jjjfj}t|dkrH|\din |\|jrd\}}}}tt rt d j j}}tt r j j}}tjtjij_fdd}jjjjd jrnjdk rntjjjfj |W5QRXjjjtj !t"j#j$jj%n.|jjtj !t"j#j$|jr |dk rt ||d jjtj<t jjtj&||d jjtj&<|dk r t ||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. Nz.Must provide batch data for current iteration.rhr,)NNNNzgWill convert to PyTorch Tensor if using compile, and casting back to MetaTensor after the forward pass.cs`jjfjjtj<tj jjtj jjtj <tj dS)N) rWrQr0outputKeysPRED fire_eventr FORWARD_COMPLETEDrSmeanLOSSLOSS_COMPLETEDr,argsroinputskwargstargetsr,r-_compute_pred_losss  $z8SupervisedTrainer._iteration.._compute_pred_loss set_to_none)metaapplied_operations)( ValueErrorrUr0rOrTralenrc isinstancer rkrl as_tensorrrrrIMAGELABELrqrQtrainrR zero_gradr`r"r&r#r$autocastrbscalerwbackwardrtr BACKWARD_COMPLETEDstepupdatersMODEL_COMPLETED) r)rorpbatch inputs_meta targets_metainputs_applied_operationstargets_applied_operationsr~r,ryr- _iterationst                zSupervisedTrainer._iteration) r?r@rArBrrrjrrCr,r,r*r-rQs,?F?cseZdZdZddddddeedddddedddddfddd d d d d d d d dddddd d dddddd ddddddfdd ZddddddZZ S)rav Generative adversarial network training based on Goodfellow et al. 2014 https://arxiv.org/abs/1406.266, inherits from ``Trainer`` and ``Workflow``. Training Loop: for each batch of data size `m` 1. Generate `m` fakes from random latent codes. 2. Update discriminator with these fakes and current batch reals, repeated d_train_steps times. 3. If g_update_latents, generate `m` fakes from new random latent codes. 4. Update generator with these fakes using discriminator feedback. 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 optimizer. d_network: discriminator (D) network architecture. d_optimizer: D optimizer function. d_loss_function: D loss function for optimizer. epoch_length: number of iterations for one epoch, default to `len(train_data_loader)`. g_inferer: inference method to execute G model forward. Defaults to ``SimpleInferer()``. d_inferer: inference method to execute D model forward. Defaults to ``SimpleInferer()``. d_train_steps: number of times to update D with real data minibatch. Defaults to ``1``. latent_shape: size of G input latent code. Defaults to ``64``. 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. d_prepare_batch: callback function to prepare batchdata for D inferer. Defaults to return ``GanKeys.REALS`` in batchdata dict. for more details please refer to: https://pytorch.org/ignite/generated/ignite.engine.create_supervised_trainer.html. g_prepare_batch: callback function to create batch of latent input for G inferer. Defaults to return random latents. for more details please refer to: https://pytorch.org/ignite/generated/ignite.engine.create_supervised_trainer.html. g_update_latents: Calculate G loss with new latent codes. Defaults to ``True``. 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. 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. 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.cuda.amp.autocast()` API, for more details: https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. Nri@FTrDrEr rFrrrGrJrHrIrKrLrMrN)rOr7rP g_network g_optimizerg_loss_function d_network d_optimizerd_loss_functionr9 g_inferer d_inferer d_train_steps latent_shaperTd_prepare_batchg_prepare_batchg_update_latentsrVrXrYrZr[r\r_r`rarbcst|tstdtj|||| |||||||||||d||_||_||_| dkr\tn| |_ ||_ ||_ | |_ | dkrtn| |_ | |_||_||_||_||_dS)Nz-train_data_loader must be PyTorch DataLoader.)rOr7rer9rTrUrVrfrZr[rgrXr_rarb)rr rr'rjrrrrrrrrrrrrrr`)r)rOr7rPrrrrrrr9rrrrrTrrrrVrXrYrZr[r\r_r`rarbr*r,r-rj]s@ zGanTrainer.__init__zdict | Sequencez,dict[str, torch.Tensor | int | float | bool]rnc CsF|dkrtd|j||jj|jf|j}|jj}|jf||j |jj|jd|j}| ||j }t d}t|jD]>}|jj|jd|||} | |j|| 7}qx|jr|jf||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_sizerOrTrir) rrUr0rOrTrare batch_sizerrrrr#zerosrangerrrr`rrritemrrrrREALSFAKESLATENTSGLOSSDLOSS) r)rorpd_inputrZg_inputZg_outputZ d_total_loss_Zdlossg_lossr,r,r-rsZ     zGanTrainer._iteration) r?r@rArBrr rrjrrCr,r,r*r-rs,HJB)4 __future__rrktypingrrrrrr#Ztorch.optim.optimizerrtorch.utils.datar monai.configr monai.datar monai.engines.utilsr r rrmonai.engines.workflowrmonai.inferersrrmonai.transformsr monai.utilsrrrmonai.utils.enumsrrrrr.monai.utils.moduler ignite.enginerrignite.metricsrOPT_IMPORT_VERSIONr__all__rrrr,r,r,r- s4           &O