U vPho- @srddlmZddlmZmZddlmZmZddlm Z m Z m Z ddl Z ddl mZddlmZddlmZmZmZdd lmZmZe rdd lmZned ejed d d\ZZdddddddddg ZGdddeZd3dddddZd4ddd d!d"d#d$dZGd%ddeZ Gd&dde Z!Gd'dde Z"d5d(d(dd d!d)d*d+dZ#d!d!d,d-d.d/dZ$d0d0d d1d2dZ%dS)6) annotations)ABCabstractmethod)CallableSequence) TYPE_CHECKINGAnycastN) IgniteInfo)apply_transform) ensure_tuple min_versionoptional_import) CommonKeysGanKeys) EventEnumz ignite.enginerbase)as_typeIterationEventsget_devices_specdefault_prepare_batch PrepareBatchPrepareBatchDefaultPrepareBatchExtraInputdefault_make_latentengine_apply_transformdefault_metric_cmp_fnc@s(eZdZdZdZdZdZdZdZdZ dS) ra Additional Events engine can register and trigger in the iteration process. Refer to the example in ignite: https://pytorch.org/ignite/generated/ignite.engine.events.EventEnum.html. These Events can be triggered during training iteration: `FORWARD_COMPLETED` is the Event when `network(image, label)` completed. `LOSS_COMPLETED` is the Event when `loss(pred, label)` completed. `BACKWARD_COMPLETED` is the Event when `loss.backward()` completed. `MODEL_COMPLETED` is the Event when all the model related operations completed. `INNER_ITERATION_STARTED` is the Event when the iteration has an inner loop and the loop is started. `INNER_ITERATION_COMPLETED` is the Event when the iteration has an inner loop and the loop is completed. Zforward_completedZloss_completedZbackward_completedZmodel_completedZinner_iteration_startedZinner_iteration_completedN) __name__ __module__ __qualname____doc__FORWARD_COMPLETEDZLOSS_COMPLETEDZBACKWARD_COMPLETEDMODEL_COMPLETEDZINNER_ITERATION_STARTEDZINNER_ITERATION_COMPLETEDr#r#H/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/engines/utils.pyr-s z#Sequence[torch.device | str] | Nonezlist[torch.device])devicesreturncCsb|dkr.ddttjD}|sPtdn"t|dkrHtdg}nt|}dd|D}|S)a Get a valid specification for one or more devices. If `devices` is None get devices for all CUDA devices available. If `devices` is and zero-length structure a single CPU compute device is returned. In any other cases `devices` is returned unchanged. Args: devices: list of devices to request, None for all GPU devices, [] for CPU. Raises: RuntimeError: When all GPUs are selected (``devices=None``) but no GPUs are available. Returns: list of torch.device: list of devices. NcSsg|]}td|dqS)zcuda:d)torchdevice.0r'r#r#r$ Ssz$get_devices_spec..zNo GPU devices available.rcpucSs$g|]}t|trt|n|qSr#) isinstancestrr(r)r*r#r#r$r,^s)ranger(cuda device_count RuntimeErrorlenr)list)r%r#r#r$rBs  F?dict[str, torch.Tensor] | torch.Tensor | Sequence[torch.Tensor]str | torch.device | Noneboolr7tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor batchdatar) non_blockingkwargsr&cKs t|tsxt|tjr0|jf||d|dfSt|dkrp|\}}|jf||d||jf||d|fStdt|tj tjr|tj jf||d||tj jf||d|fSt j |kr|t j jf||d|S|tj jf||d|dfS)a Default function to prepare the data for current iteration. The input `batchdata` is either a single tensor, a pair of tensors, or a dictionary of data. In the first case the return value is the tensor and None, in the second case the return value is the two tensors, and in the dictionary case the return value depends on what keys are present. if `CommonKeys.IMAGE` and `CommonKeys.LABEL` are present then the tensors they key to are returned, if only `CommonKeys.IMAGE` is present that tensor and None is returned. If `CommonKeys.REALS` is present this is returned with None. All returned tensors are moved to the given device using the given non-blocking argument before being returned. This function implements the expected API for a `prepare_batch` callable in Ignite: https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html Args: batchdata: input batch data which is either a single tensor, a pair, or a dictionary device: device to move every returned tensor to non_blocking: equivalent argument for `Tensor.to` kwargs: further arguments for `Tensor.to` Returns: image, label(optional). r)r<NzWDefault prepare_batch expects a single tensor, a tensor pair, or dictionary input data.) r.dictr(Tensortor4AssertionErrorgetrLABELIMAGErREALS)r;r)r<r=imagelabelr#r#r$rbs     c@s,eZdZdZed ddddddd d ZdS) ra Interface of customized prepare_batch in the trainer or evaluator workflows. It takes the data of current batch, target device and non_blocking flag as input. Args `batchdata`, `device`, `non_blocking` refer to the ignite API: https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html. `kwargs` supports other args for `Tensor.to()` API. NFdict[str, torch.Tensor]r7r8rr:cKstd|jjddS)Nz Subclass z must implement this method.)NotImplementedError __class__rselfr;r)r<r=r#r#r$__call__szPrepareBatch.__call__)NF)rrrr rrOr#r#r#r$rs c@s(eZdZdZd dddddd d d ZdS) rzo This wraps `default_prepare_batch` to return `image` and `label` only, so is consistent with its API. NFr6r7r8rr9r:cKst|||f|S)z Args `batchdata`, `device`, `non_blocking` refer to the ignite API: https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html. `kwargs` supports other args for `Tensor.to()` API. )rrMr#r#r$rOs zPrepareBatchDefault.__call__)NF)rrrr rOr#r#r#r$rsc@s8eZdZdZdddddZdd d d d d dddZdS)ra Customized prepare batch callable for trainers or evaluators which support extra input data for the network. Extra items are specified by the `extra_keys` parameter and are extracted from the input dictionary (ie. the batch). This uses `default_prepare_batch` but requires dictionary inputs. Args: extra_keys: If a string or sequence of strings is provided, values from the input dictionary are extracted from those keys and passed to the network as extra positional arguments. If a dictionary is provided, every pair `(k, v)` in that dictionary will become a new keyword argument assigning to `k` the value in the input dictionary keyed to `v`. z$str | Sequence[str] | dict[str, str]None) extra_keysr&cCs ||_dS)N)rQ)rNrQr#r#r$__init__szPrepareBatchExtraInput.__init__NFrJr7r8rz.tuple[torch.Tensor, torch.Tensor, tuple, dict]r:c  stf\}}t}t}dddfdd } t|jtttfrlt|jD]} || | qVn2t|jtr|j D]\} } | | | | iqt t j |t t j |t||fS)z Args `batchdata`, `device`, `non_blocking` refer to the ignite API: https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html. `kwargs` supports other args for `Tensor.to()` API. r/ torch.Tensor)keyr&cs.|}t|tjr*|jfd}|SNr>)r.r(rArB)rTdatar;r)r=r<r#r$ _get_datas z2PrepareBatchExtraInput.__call__.._get_data)rr5r@r.rQr/tupler appenditemsupdater r(rA) rNr;r)r<r=rHrIargs_kwargs_rXkvr#rWr$rOs  zPrepareBatchExtraInput.__call__)NF)rrrr rRrOr#r#r#r$rs  intrS) num_latents latent_sizer)r<r=r&cKst||jf||d|SrU)r(randnrB)rbrcr)r<r=r#r#r$rszCallable[..., dict]ztuple[Any, Any])batchoutput transformr&cCst|trxt|trxt|}||t||}t|tsBtd|D]*\}}||ksb||krl|||<qJ|||<qJn t||}||fS)z Apply transform on `batch` and `output`. If `batch` and `output` are dictionaries, temporarily combine them for the transform, otherwise, apply the transform for `output` data only. zIWith a dict supplied to apply_transform a single dict return is expected.)r.r@r\r rCr[)rerfrgrVZtransformed_datar_r`r#r#r$rs      float)current_metric prev_bestr&cCs||kS)a The default function to compare metric values between current metric and previous best metric. Args: current_metric: metric value of current round computation. prev_best: the best metric value of previous rounds to compare with. r#)rirjr#r#r$rs )N)NF)NF)& __future__rabcrrcollections.abcrrtypingrrr r( monai.configr monai.transformsr monai.utilsr r rmonai.utils.enumsrr ignite.enginerOPT_IMPORT_VERSION___all__rrrrrrrrrr#r#r#r$ sP     "45