U Ph!@s,ddlmZddlZddlZddlmZddlmZejdkrJddl m Z ddl m Z ddl Z ddl mZddlmZdd lmZmZed ejed \ZZd d ddgZdd Ze dddddd Ze dddddd Ze dddddd Zd#dddddd Zd$dddd d!dZGd"ddeZdS)%) annotationsN)Callable)Filter))Literal)overload) IgniteInfo) min_versionoptional_importZignite distributedget_dist_deviceevenly_divisible_all_gatherstring_list_all_gather RankFiltercCsNtrJt}|dkr8tjr8tdtjS|dkrJtdSdS)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)backendrE/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/utils/dist.pyr #s  torch.Tensorz Literal[True])dataconcatreturncCsdSNrr r!rrrr4szLiteral[False]list[torch.Tensor]cCsdSr#rr$rrrr8sboolz!torch.Tensor | list[torch.Tensor]cCsdSr#rr$rrrr<sTcst|tjstd|dkr.|jdnddddfdd }dddfd d }tr~tdkrr|S||d }n0t rt rt dkr|S||d }n|S|rtj |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.rrr%)r r"cst}j|dkr(dntjg|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)rcsg|]}tqSrr zeros_like.0_) length_tensorrr cszJevenly_divisible_all_gather.._torch_all_gather..cSsg|]}t|qSr)intitem)r+irrrr.esr'Ndimcsg|]}tqSrr(r*r rrr.lscs8g|]0\}}dkr|dn|d|dfqS)rN.)squeezeto)r+ol)ndims orig_devicerrr.os)r rr6 unsqueezer as_tensorrangerget_world_size all_gathermaxlistshapecatnew_fullzip)r rall_lensZ all_lens_max_lensizeoutputlengthr9)r r-r:r_torch_all_gatherWs   z6evenly_divisible_all_gather.._torch_all_gathercsdkr|dn|}t}t|krfgt|jdd}tj|||dgdd}t|dkrttj ddSfddt |DS)zi Implementation based on PyTorch ignite package, it can support more kinds of backends. rr'Nr2cs,g|]$\}}|||dfqS).r)r+r1r8rGrIrrr.szKevenly_divisible_all_gather.._ignite_all_gather..) r;idistr?r@rArBrrCrDunbind enumerate)r rFrHrJrMr_ignite_all_gatherqs  z7evenly_divisible_all_gather.._ignite_all_gatherr4r2) isinstancerTensor ValueError ndimensionrB has_igniterNr>rrrrC)r r!rLrQrIrrJrr@s       z list[str]str)strings delimiterr"cs~d}trt}ntr*tr*t}|dkr6|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)r!cs$g|]}t|dqS)r[) bytearraytolistdecodesplit)r+grZrrr.sz*string_list_all_gather..cSsg|]}|D]}|q qSrr)r+kr1rrrr.s) rVrNr>rrrjoinrrtensorr]long)rYrZ world_sizejoinedgatheredZ _gatheredrrbrrs  cs:eZdZdZdddfdddfdd 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)rankrrrzRankFilter.z int | Noner)rj filter_fncsdt||_tr8tr8|dk r,|nt|_n(tj rZtj dkrZt 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__rmrrrget_rankrjrr device_countwarningswarn)selfrjrm __class__rrros zRankFilter.__init__cGs ||jSr#)rmrj)rt_argsrrrfilterszRankFilter.filter)__name__ __module__ __qualname____doc__rorx __classcell__rrrurrs )T)rW) __future__rsysrrcollections.abcrloggingr version_infotypingrrrtorch.distributedr rZ monai.configr monai.utils.moduler r ZOPT_IMPORT_VERSIONrNrV__all__r rrrrrrr s.         T