o  i!@sddlmZddlZddlmZddlmZddlmZm Z ddl Z ddl m Z ddlmZddlmZmZedejed \ZZgd Zd d Ze d%ddZe d&ddZe d'ddZd(d'ddZd)d*d!d"ZGd#d$d$eZdS)+) annotationsN)Callable)Filter)Literaloverload) IgniteInfo) min_versionoptional_importZignite distributed)get_dist_deviceevenly_divisible_all_gatherstring_list_all_gather RankFiltercCsNtr%t}|dkrtjrtdtjS|dkr%tdSdS)a Get the expected target device in the native PyTorch distributed data parallel. For NCCL backend, return GPU device of current process. For GLOO backend, return CPU. For any other backends, return None as the default, tensor.to(None) will not change the device. ncclzcuda:gloocpuN)distis_initialized get_backendtorchcuda is_availabledevicecurrent_device)backendrR/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/utils/dist.pyr s r data torch.Tensorconcat Literal[True]returncCdSNrrrrrrr /r Literal[False]list[torch.Tensor]cCr"r#rr$rrrr 3r%bool!torch.Tensor | list[torch.Tensor]cCr"r#rr$rrrr 7r%Tcst|tjs td|dkr|jdnddfdd }dfd d }tr9tdkr3|S||d }nt rOt rOt dkrI|S||d }n|S|rZtj |dd S|S)a) Utility function for distributed data parallel to pad at first dim to make it evenly divisible and all_gather. The input data of every rank should have the same number of dimensions, only the first dim can be different. Note: If has ignite installed, will execute based on ignite distributed APIs, otherwise, if the native PyTorch distributed group initialized, will execute based on native PyTorch distributed APIs. Args: data: source tensor to pad and execute all_gather in distributed data parallel. concat: whether to concat the gathered list to be a Tensor, if False, return a list of Tensors, similar behavior as torch.distributed.all_gather(). default to True. Note: The input data on different ranks must have exactly same `dtype`. z"input data must be PyTorch Tensor.rrrr!r'cst}j|dkrdntjg|dfddttD}t |dd|D}t |}|kr[|gt j dd}tj |dgddfd dttD}t |fd dt||DS) zY Implementation based on native PyTorch distributed data parallel APIs. r)rcg|]}tqSrr zeros_like.0_) length_tensorrr ^zJevenly_divisible_all_gather.._torch_all_gather..cSsg|]}t|qSr)intitem)r/irrrr2`sr*Ndimcr+rr,r.rrrr2gr3cs8g|]\}}dkr|dn|d|dfqS)rN.)squeezeto)r/ol)ndims orig_devicerrr2js8)r rr; unsqueezer as_tensorrangerget_world_size all_gathermaxlistshapecatnew_fullzip)rrall_lensZ all_lens_max_lensizeoutputlengthr>)rr1r?r_torch_all_gatherRs   z6evenly_divisible_all_gather.._torch_all_gathercsdkr |dn|}t}t|kr3gt|jdd}tj|||dgdd}t|dkrEttj ddSfddt |DS)zi Implementation based on PyTorch ignite package, it can support more kinds of backends. rr*Nr7cs,g|]\}}|||dfqS).r)r/r6r=rLrNrrr2~s,zKevenly_divisible_all_gather.._ignite_all_gather..) r@idistrDrErFrGrrHrIunbind enumerate)rrKrMrOrRr_ignite_all_gatherls  z7evenly_divisible_all_gather.._ignite_all_gatherr9r7N)rrr!r') isinstancerTensor ValueError ndimensionrG has_igniterSrCrrrrH)rrrQrVrNrrOrr ;s      strings list[str] delimiterstrcs~d}tr t}n trtrt}|dkr|S|}ttj t |dtj ddd}fdd|D}dd|DS) a Utility function for distributed data parallel to all gather a list of strings. Refer to the idea of ignite `all_gather(string)`: https://pytorch.org/ignite/v0.4.5/distributed.html#ignite.distributed.utils.all_gather. Note: If has ignite installed, will execute based on ignite distributed APIs, otherwise, if the native PyTorch distributed group initialized, will execute based on native PyTorch distributed APIs. Args: strings: a list of strings to all gather. delimiter: use the delimiter to join the string list to be a long string, then all gather across ranks and split to a list. default to " ". r*utf-8)dtypeF)rcs$g|]}t|dqS)ra) bytearraytolistdecodesplit)r/gr_rrr2s$z*string_list_all_gather..cSsg|] }|D]}|qqSrr)r/kr6rrrr2s) r[rSrCrrrjoinr rtensorrclong)r]r_ world_sizejoinedgatheredZ _gatheredrrhrr s  r cs4eZdZdZdddfd fd d Zd d ZZS)raW The RankFilter class is a convenient filter that extends the Filter class in the Python logging module. The purpose is to control which log records are processed based on the rank in a distributed environment. Args: rank: the rank of the process in the torch.distributed. Default is None and then it will use dist.get_rank(). filter_fn: an optional lambda function used as the filtering criteria. The default function logs only if the rank of the process is 0, but the user can define their own function to implement custom filtering logic. NcCs|dkS)Nrr)rankrrrszRankFilter.rp int | None filter_fnrcslt||_tr tr |dur||_dSt|_dStj r1tj dkr1t dd|_dS)Nr*zThe torch.distributed is either unavailable and uninitiated when RankFilter is instantiated. If torch.distributed is used, please ensure that the RankFilter() is called after torch.distributed.init_process_group() in the script. r) super__init__rsrrrget_rankrprr device_countwarningswarn)selfrprs __class__rrrus   zRankFilter.__init__cGs ||jSr#)rsrp)rz_argsrrrfilters zRankFilter.filter)rprrrsr)__name__ __module__ __qualname____doc__rur~ __classcell__rrr{rrs r)rrrr r!r)rrrr&r!r')rrrr(r!r))T)r\)r]r^r_r`r!r^) __future__rrxcollections.abcrloggingrtypingrrrtorch.distributedr rmonai.utils.enumsrmonai.utils.modulerr ZOPT_IMPORT_VERSIONrSr[__all__r r r rrrrrs*          T