o  i*=@sbddlmZddlmZmZddlmZmZmZddl m Z m Z m Z ddl Z ddlmZddlmZddlmZmZmZmZddlmZmZe rOdd lmZn ed ejed d d \ZZgdZGdddeZd?d@ddZ  dAdBd!d"Z!Gd#d$d$eZ"Gd%d&d&e"Z#Gd'd(d(e"Z$Gd)d*d*e"Z%Gd+d,d,e%Z&  dAdCd1d2Z'dDd8d9Z(dEd=d>Z)dS)F) annotations)ABCabstractmethod)CallableMappingSequence) TYPE_CHECKINGAnycastN)apply_transform) IgniteInfo ensure_tuple min_versionoptional_import) CommonKeysGanKeys) EventEnumz ignite.enginerbase)as_type) IterationEventsget_devices_specdefault_prepare_batch PrepareBatchPrepareBatchDefaultPrepareBatchExtraInputDiffusionPrepareBatchVPredictionPrepareBatchdefault_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. forward_completedZloss_completedZbackward_completedZmodel_completedZinner_iteration_startedZinner_iteration_completedN) __name__ __module__ __qualname____doc__FORWARD_COMPLETEDZLOSS_COMPLETEDZBACKWARD_COMPLETEDMODEL_COMPLETEDZINNER_ITERATION_STARTEDZINNER_ITERATION_COMPLETEDr'r'U/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/engines/utils.pyr/s rdevices#Sequence[torch.device | str] | Nonereturnlist[torch.device]cCsb|durddttjD}|stdnt|dkr$tdg}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( Usz$get_devices_spec..zNo GPU devices available.rcpucSs$g|]}t|trt|n|qSr') isinstancestrr.r/r0r'r'r(r2`s$)ranger.cuda device_count RuntimeErrorlenr/list)r)r'r'r(rDs rF batchdata?dict[str, torch.Tensor] | torch.Tensor | Sequence[torch.Tensor]r/str | torch.device | None non_blockingboolkwargsr 7tuple[torch.Tensor, torch.Tensor | None] | torch.TensorcKs t|tsr?r@rAr r+cKstd|jjd)Nz Subclass z must implement this method.)NotImplementedError __class__r!selfr<r/r?rAr'r'r(__call__szPrepareBatch.__call__NF) r<rOr/r>r?r@rAr r+r )r!r"r#r$rrTr'r'r'r(rs rc@s eZdZdZ  ddddZdS)rzo This wraps `default_prepare_batch` to return `image` and `label` only, so is consistent with its API. NFr<r=r/r>r?r@rAr r+rBcKst|||fi|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. )rrRr'r'r(rTs zPrepareBatchDefault.__call__rU r<r=r/r>r?r@rAr r+rB)r!r"r#r$rTr'r'r'r(rs rc@s*eZdZdZdddZ  ddddZdS)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`. extra_keys$str | Sequence[str] | dict[str, str]r+NonecCs ||_dSN)rW)rSrWr'r'r(__init__s zPrepareBatchExtraInput.__init__NFr<rOr/r>r?r@rAr .tuple[torch.Tensor, torch.Tensor, tuple, dict]c  stfi\}}t}t}dfdd } t|jtttfr5t|jD] } || | q*nt|jtrN|j D] \} } | | | | iq@t 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. keyr5r+ torch.Tensorcs.|}t|tjr|jdd}|SNrCr')r4r.rFrG)r]datar<r/rAr?r'r( _get_datas z2PrepareBatchExtraInput.__call__.._get_dataN)r]r5r+r^)rr;rEr4rWr5tupler appenditemsupdater r.rF) rSr<r/r?rArMrNargs_kwargs_rbkvr'rar(rTs   zPrepareBatchExtraInput.__call__)rWrXr+rYrU r<rOr/r>r?r@rAr r+r\)r!r"r#r$r[rTr'r'r'r(rs  rc@sJeZdZdZd!d"d d Zd#d dZd#ddZd$ddZ  d%d&dd ZdS)'ra This class is used as a callable for the `prepare_batch` parameter of engine classes for diffusion training. Assuming a supervised training process, it will generate a noise field using `get_noise` for an input image, and return the image and noise field as the image/target pair plus the noise field the kwargs under the key "noise". This assumes the inferer being used in conjunction with this class expects a "noise" parameter to be provided. If the `condition_name` is provided, this must refer to a key in the input dictionary containing the condition field to be passed to the inferer. This will appear in the keyword arguments under the key "condition". Nnum_train_timestepsintcondition_name str | Noner+rYcCs||_||_dSrZ)rnrl)rSrlrnr'r'r(r[s zDiffusionPrepareBatch.__init__imagesr^cCs t|S)zdReturns the noise tensor for input tensor `images`, override this for different noise distributions.)r. randn_likerSrpr'r'r( get_noises zDiffusionPrepareBatch.get_noisecCs"tjd|j|jdf|jdS)z]Get a timestep, by default this is a random integer between 0 and `self.num_train_timesteps`.r)r/)r.randintrlshaper/longrrr'r'r( get_timestepss"z#DiffusionPrepareBatch.get_timestepsnoise timestepscCs|S)zNReturn the target for the loss function, this is the `noise` value by default.r'rSrprxryr'r'r( get_target sz DiffusionPrepareBatch.get_targetFr<rOr/r>r?r@rAr r\c Kst|||fi|\}}||j|fd|i|}||j|fd|i|}||||j|fd|i|} ||d} |jdurWt|trW||jj|fd|i|| d<|| d| fS)Nr?)rxry conditionr')rrsrGrwr{rnr4r) rSr<r/r?rArp_rxrytargetZ infer_kwargsr'r'r(rTs    zDiffusionPrepareBatch.__call__rZ)rlrmrnror+rY)rpr^r+r^)rpr^rxr^ryr^r+r^rUrk) r!r"r#r$r[rsrwr{rTr'r'r'r(rs   rcs,eZdZdZddfd d Zd dZZS)ra  This class is used as a callable for the `prepare_batch` parameter of engine classes for diffusion training. Assuming a supervised training process, it will generate a noise field using `get_noise` for an input image, and from this compute the velocity using the provided scheduler. This value is used as the target in place of the noise field itself although the noise is field is in the kwargs under the key "noise". This assumes the inferer being used in conjunction with this class expects a "noise" parameter to be provided. If the `condition_name` is provided, this must refer to a key in the input dictionary containing the condition field to be passed to the inferer. This will appear in the keyword arguments under the key "condition". N scheduler nn.Modulerlrmrnror+rYcstj||d||_dS)N)rlrn)superr[r)rSrrlrnrQr'r(r[3s z VPredictionPrepareBatch.__init__cCs|j|||SrZ)rZ get_velocityrzr'r'r(r{7sz"VPredictionPrepareBatch.get_targetrZ)rrrlrmrnror+rY)r!r"r#r$r[r{ __classcell__r'r'rr(r%s r num_latentsrm latent_sizer^cKst||jd||d|Sr_)r.randnrG)rrr/r?rAr'r'r(r;srbatchoutput transformCallable[..., dict]tuple[Any, Any]cCst|tr?t|tr?t|}||t||}t|ts!td|D]\}}||vs1||vr6|||<q%|||<q%||fSt||}||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.)r4rErfr rHre)rrrr`Ztransformed_datarirjr'r'r(rEs      rcurrent_metricfloat prev_bestcCs||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')rrr'r'r(ras rrZ)r)r*r+r,rUrV) rrmrrmr/r>r?r@rAr r+r^)rr rr rrr+r)rrrrr+r@)* __future__rabcrrcollections.abcrrrtypingrr r r.torch.nnnnmonai.transformsr monai.utilsr r rrmonai.utils.enumsrr ignite.enginerOPT_IMPORT_VERSIONr}__all__rrrrrrrrrrrr'r'r'r(s<      " 422