o , i@sddlmZddlZddlZddlmZddlZddlZddl m Z ddl m m Z ddlmZmZmZmZddlmZddlmZmZddlmZedd d \ZZed d d \ZZgd ZejjGdddej j Z!ejjGdddej j Z"Gddde"Z#ddZ$Gddde j%Z&GdddeZ'GdddeZ(GdddeZ)GdddeZ*Gd d!d!e j Z+Gd"d#d#e!Z,Gd$d%d%e j Z-Gd&d'd'e j Z.Gd(d)d)e.Z/Gd*d+d+e.Z0dS),) annotationsN)Optional)ActiConvNormBlockFactorizedIncreaseBlockFactorizedReduceBlockP3DActiConvNormBlock)Conv) get_act_layerget_norm_layer)optional_importz scipy.sparse csr_matrixnamezscipy.sparse.csgraphdijkstra)DiNTSTopologyConstructionTopologyInstanceTopologySearchc@seZdZdZd ddZd S) CellInterfacez"interface for torchscriptable Cellx torch.TensorweightOptional[torch.Tensor]returncCdSNselfrrrr[/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/dints.pyforward,zCellInterface.forwardNrrrrrr__name__ __module__ __qualname____doc__r rrrrr(rc@seZdZdZdddZdS) StemInterfacez"interface for torchscriptable StemrrrcCrrrrrrrrr 4r!zStemInterface.forwardNrrrrr#rrrrr)0r(r)cs*eZdZdZfddZd ddZZS) StemTSz wrapper for torchscriptable Stemcsttjj||_dSr)super__init__torchnn Sequentialmod)rr2 __class__rrr.;s zStemTS.__init__rrrcCs ||Sr)r2r*rrrr ?s zStemTS.forwardr+r$r%r&r'r.r __classcell__rrr3rr,8s r,cCs>||kr dgdggSt|d|}dd|Ddd|DS)z>use depth first search to find all path activation combinationrcSg|]}dg|qSrr.0_rrr Hz_dfs..cSr8)r7rr:rrrr=Hr>)_dfs)nodepathschildrrrr?Cs r?cseZdZfddZZS)_IdentityWithRAMCostcstj|i|d|_dSNrr-r.ram_cost)rargskwargsr3rrr.Ms z_IdentityWithRAMCost.__init__r$r%r&r.r6rrr3rrCKsrCcs2eZdZdZdddddiffdfdd ZZS)_ActiConvNormBlockWithRAMCosta!The class wraps monai layers with ram estimation. The ram_cost = total_ram/output_size is estimated. Here is the estimation: feature_size = output_size/out_channel total_ram = ram_cost * output_size total_ram = in_channel * feature_size (activation map) + in_channel * feature_size (convolution map) + out_channel * feature_size (normalization) = (2*in_channel + out_channel) * output_size/out_channel ram_cost = total_ram/output_size = 2 * in_channel/out_channel + 1 RELUINSTANCEaffineT in_channelint out_channel kernel_sizepadding spatial_dimsact_name tuple | str norm_namec s.t|||||||d||d|_dS)Nr7rE)rrOrQrRrSrTrUrWr3rrr.^s z&_ActiConvNormBlockWithRAMCost.__init__)rOrPrQrPrRrPrSrPrTrPrUrVrWrV)r$r%r&r'r.r6rrr3rrJRs  rJcs.eZdZdddddiffdfdd ZZS) _P3DActiConvNormBlockWithRAMCostrrLrMrNTrOrPrQrRrSp3dmoderUrVrWc s.t|||||||dd|||_dSNrXrE)rrOrQrRrSrZrUrWr3rrr.ns z)_P3DActiConvNormBlockWithRAMCost.__init__)rOrPrQrPrRrPrSrPrZrPrUrVrWrVrIrrr3rrYls  rYc.eZdZdddddiffdfd d ZZS)#_FactorizedIncreaseBlockWithRAMCostrKrLrMrNTrOrPrQrTrUrVrWcs*t|||||d||d|_dSr[rErrOrQrTrUrWr3rrr.sz,_FactorizedIncreaseBlockWithRAMCost.__init__ rOrPrQrPrTrPrUrVrWrVrIrrr3rr]~  r]cr\)!_FactorizedReduceBlockWithRAMCostrKrLrMrNTrOrPrQrTrUrVrWcs0t|||||||d|jd|_dS)NrXrK)r-r. _spatial_dimsrFr^r3rrr.sz*_FactorizedReduceBlockWithRAMCost.__init__r_rIrrr3rrar`racs0eZdZdZddfdd Zddd dZZS)MixedOpa# The weighted averaging of cell operations. Args: c: number of output channels. ops: a dictionary of operations. See also: ``Cell.OPS2D`` or ``Cell.OPS3D``. arch_code_c: binary cell operation code. It represents the operation results added to the output. NcrPopsdictcs^t|durtt|}t|_t||D]\}}|dkr,|j |||qdSrD) r-r.nponeslenr0 ModuleListrezipappend)rrdre arch_code_cZarch_cop_namer3rrr.s  zMixedOp.__init__rrrrcCsVd}|dur ||}t|jD]\}}|dur|||n |||||}q|S)z Args: x: input tensor. weight: learnable architecture weights for cell operations. arch_code_c are derived from it. Return: out: weighted average of the operation results. N)to enumeratere)rrroutidx_oprrrr s  *zMixedOp.forwardr)rdrPrerf)rrrrr5rrr3rrcs rccseZdZdZdZdddddZddddd dd dd dd Zeee e d Z ddddddiffd$fdd Z d%d"d#Z ZS)&Cella The basic class for cell operation search, which contains a preprocessing operation and a mixed cell operation. Each cell is defined on a `path` in the topology search space. Args: c_prev: number of input channels c: number of output channels rate: resolution change rate. It represents the preprocessing operation before the mixed cell operation. ``-1`` for 2x downsample, ``1`` for 2x upsample, ``0`` for no change of resolution. arch_code_c: cell operation code rKcCtSrrC_crrrz Cell.cCt||ddddS)NrKr7rXrSrTrJrdrrrrz skip_connectZconv_3x3cCrvrrwrxrrrrzr{cCst||ddddS)NrKr7r}r~rrrrrzrcCr|)NrKr7rrSrZrYrrrrrzrcCst||ddddS)NrKr7rrrrrrrzrcCr|)NrKr7rXrrrrrrrzrrZ conv_3x3x3Z conv_3x3x1Z conv_3x1x3Z conv_1x3x3)updownidentityalign_channelsNrLrMrNTc_prevrPrdraterTrUrVrWc sNt|_|_|_|dkr#jd||jjjd_n3|dkr8jd||jjjd_n||krDjd_njd||ddjjjd_d d fd d d _d d fdd fdd fdd fdd d_i_ jdkrj_ njdkrj_ n t djdt |j |_ dS)Nr)rTrUrWr7rrrrcSrvrrwrxrrrrzr{zCell.__init__..c t||dddjjdS)NrKr7rXrSrTrUrWrJ _act_name _norm_namerrrrrzrcSrvrrwrxrrrrzr{c st||dddjjdS)NrKr7rrrrrrrzrc r)NrKr7rrSrZrUrWrYrrrrrrrzrc st||dddjjdS)NrKr7rrrrrrrz rc r)NrKr7rXrrrrrrrz#rrrXrKSpatial dimensions  is not supported.) r-r.rbrrConnOPS preprocessOPS2DOPS3DZOPSNotImplementedErrorrcop)rrrdrrmrTrUrWr3rrr.sB            z Cell.__init__rrrrrcCs||}|||}|S)zi Args: x: input tensor weight: weights for different operations. )rrrrrrr 2s  z Cell.forward) rrPrdrPrrPrTrPrUrVrWrVr")r$r%r&r' DIRECTIONSrrr]rarCrJrr.r r6rrr3rrus.     BrucsHeZdZdZddddifdddfdfdd ZddZdddZZS)ra Reimplementation of DiNTS based on "DiNTS: Differentiable Neural Network Topology Search for 3D Medical Image Segmentation ". The model contains a pre-defined multi-resolution stem block (defined in this class) and a DiNTS space (defined in :py:class:`monai.networks.nets.TopologyInstance` and :py:class:`monai.networks.nets.TopologySearch`). The stem block is for: 1) input downsample and 2) output upsample to original size. The model downsamples the input image by 2 (if ``use_downsample=True``). The downsampled image is downsampled by [1, 2, 4, 8] times (``num_depths=4``) and used as input to the DiNTS search space (``TopologySearch``) or the DiNTS instance (``TopologyInstance``). - ``TopologyInstance`` is the final searched model. The initialization requires the searched architecture codes. - ``TopologySearch`` is a multi-path topology and cell operation search space. The architecture codes will be initialized as one. - ``TopologyConstruction`` is the parent class which constructs the instance and search space. To meet the requirements of the structure, the input size for each spatial dimension should be: divisible by 2 ** (num_depths + 1). Args: dints_space: DiNTS search space. The value should be instance of `TopologyInstance` or `TopologySearch`. in_channels: number of input image channels. num_classes: number of output segmentation classes. act_name: activation name, default to 'RELU'. norm_name: normalization used in convolution blocks. Default to `InstanceNorm`. spatial_dims: spatial 2D or 3D inputs. use_downsample: use downsample in the stem. If ``False``, the search space will be in resolution [1, 1/2, 1/4, 1/8], if ``True``, the search space will be in resolution [1/2, 1/4, 1/8, 1/16]. node_a: node activation numpy matrix. Its shape is `(num_depths, num_blocks + 1)`. +1 for multi-resolution inputs. In model searching stage, ``node_a`` can be None. In deployment stage, ``node_a`` cannot be None. rLrMrNTrKN in_channelsrP num_classesrUrVrWrTuse_downsampleboolc st||_|j|_|j|_|j|_|dvr td|d||_|dur4t |jd|jf|_ n||_ t t j |f} t |_t |_t t|jd|jd|||d| |jd|ddddddd|_|jd krpd nd } t|jD]} |rtt jdd | | dd | ||j| d ddddddt|||j| dt|d| |j| |j| dd d dddddt|||j| dd|jt| <tt|d| |j| d|j| d ddddddt|||j| dt jd | dd |jt| <qwtt jdd | | dd | ||j| d ddddddt|||j| d|jt| <tt|d| |j| |jt| ddd ddddddt|||jt| dddt jd | dk| dd |jt| <qwdS)N)rXrKrrr7r)rUrWrTT)r out_channelsrRstriderSgroupsbiasdilationrK trilinearbilinearrX) scale_factormode align_cornersF)rrTchannelsr )r-r. dints_space filter_nums num_blocks num_depthsrrbr/rhnode_arCONVr0 ModuleDict stem_downstem_upr1r stem_finalsranger,Upsampler r strmax) rrrrrUrWrTrr conv_typerres_idxr3rrr.cs          zDiNTS.__init__cCsdd|DS)NcSsg|]\}}|qSrrr;rparamrrrr=sz+DiNTS.weight_parameters..named_parametersrrrrweight_parametersszDiNTS.weight_parametersrrc Csg}t|jD]#}|jt|}||}|jd|r"||q|t|q| |}|j d}d}t d} t|jdddD]'} |j t| } |rZ| || | } qE|j|d| rld}| || } qE| | } | S)zd Prediction based on dynamic arch_code. Args: x: input tensor. rr7FrT)rrrrr rrlr/ zeros_likerremptyrr) rrinputsdZ_mod_wx_outoutputsblk_idxstart_temprZ_mod_up predictionrrrr s*      z DiNTS.forward) rrPrrPrUrVrWrVrTrPrr)rr)r$r%r&r'r.rr r6rrr3rr=s* rc sHeZdZdZddeddddddd ifd d f dfdd ZddZZS)ra? The base class for `TopologyInstance` and `TopologySearch`. Args: arch_code: `[arch_code_a, arch_code_c]`, numpy arrays. The architecture codes defining the model. For example, for a ``num_depths=4, num_blocks=12`` search space: - `arch_code_a` is a 12x10 (10 paths) binary matrix representing if a path is activated. - `arch_code_c` is a 12x10x5 (5 operations) binary matrix representing if a cell operation is used. - `arch_code` in ``__init__()`` is used for creating the network and remove unused network blocks. If None, all paths and cells operations will be used, and must be in the searching stage (is_search=True). channel_mul: adjust intermediate channel number, default is 1. cell: operation of each node. num_blocks: number of blocks (depth in the horizontal direction) of the DiNTS search space. num_depths: number of image resolutions of the DiNTS search space: 1, 1/2, 1/4 ... in each dimension. use_downsample: use downsample in the stem. If False, the search space will be in resolution [1, 1/2, 1/4, 1/8], if True, the search space will be in resolution [1/2, 1/4, 1/8, 1/16]. device: `'cpu'`, `'cuda'`, or device ID. Predefined variables: `filter_nums`: default to 32. Double the number of channels after downsample. topology related variables: - `arch_code2in`: path activation to its incoming node index (resolution). For depth = 4, arch_code2in = [0, 1, 0, 1, 2, 1, 2, 3, 2, 3]. The first path outputs from node 0 (top resolution), the second path outputs from node 1 (second resolution in the search space), the third path outputs from node 0, etc. - `arch_code2ops`: path activation to operations of upsample 1, keep 0, downsample -1. For depth = 4, arch_code2ops = [0, 1, -1, 0, 1, -1, 0, 1, -1, 0]. The first path does not change resolution, the second path perform upsample, the third perform downsample, etc. - `arch_code2out`: path activation to its output node index. For depth = 4, arch_code2out = [0, 0, 1, 1, 1, 2, 2, 2, 3, 3], the first and second paths connects to node 0 (top resolution), the 3,4,5 paths connects to node 1, etc. N?rKrLrMrNTcpu arch_code list | None channel_mulfloatrrPrrTrUrVrWrrdevicerc  sttddt|dD} fdd| D|_||_||_tdt j d|t | ||_ ||_ ||_| |_| |_d|_|j dkrQt|j|_n |j dkr\t|j|_gg} } ttj|jdD]}| |dtjd|dtjqkgd |jdd }t|jD] }| |||gq| dd } | |_||_| |_|durt|jt|jf|j}t|jt|j|jf|j}nt|d|j}t !t|dtj"|j|j}||_#||_$t%&|_'t|jD]G}tt|jD]=}|j#||fdkrD||j|j|t | |j|j|t | |j||j$||f|j |j |j|j't(||f<qqdS) NcSsg|]}dd|qS) rXr)r;_irrrr=;z1TopologyConstruction.__init__..r7csg|]}t|qSr)rP)r;Zn_featrrrr=<rzC{} - Length of input patch is recommended to be a multiple of {:d}.rXrrK)rrr7r))r-r.tuplerrrrprintformatdatetimenowrPrbrrrr num_cell_opsrirrrurrlextend arch_code2in arch_code2ops arch_code2outr/rhrp from_numpyFone_hotint64 arch_code_armr0r cell_treer)rrrcellrrrTrUrWrrZn_featsrrirmrrmrrr3rrr.,sj     ( $(  zTopologyConstruction.__init__cCsdS)zOThis function to be implemented by the architecture instances or search spaces.Nrr*rrrr vszTopologyConstruction.forward)rrrrrrPrrPrTrPrUrVrWrVrrrrr$r%r&r'rur.r r6rrr3rrs' Jrc sJeZdZdZddeddddddd ifd d f dfdd Zd ddZZS)!rz` Instance of the final searched architecture. Only used in re-training/inference stage. NrrrKrLrMrNTrrrrrPrrTrUrVrWrrrrc s6|dur tdtj||||||||| | d dS)Q Initialize DiNTS topology search space of neural architectures. Nz*arch_code not provided when not searching. rrrrrrTrUrWrr)warningswarnr-r.) rrrrrrrTrUrWrrr3rrr.s  zTopologyInstance.__init__rlist[torch.Tensor]rc Cs|}t|jD]H}tjd|dj|djdg|j}t|j|j D])\}}|rL|j t ||f}|j ||j |dd}||j||||j|<q#|}q|S)z4 Args: x: input tensor. rordtyperN)rr)rrr/tensorrrrrqrdatarrr rr) rrrrrr activationr2_outrrrr s$zTopologyInstance.forward)rrrrPrrPrTrPrUrVrWrVrrrr)rrrrrrrr3rr{s  rc seZdZUdZded<ded<deddddd d d d ifd d f d1fdd Zd2d d!Zd"d#Zd3d4d&d'Z d3d5d)d*Z d+d,Z d-d.Z d/d0Z ZS)6ra DiNTS topology search space of neural architectures. Examples: .. code-block:: python from monai.networks.nets.dints import TopologySearch topology_search_space = TopologySearch( channel_mul=0.5, num_blocks=8, num_depths=4, use_downsample=True, spatial_dims=3) topology_search_space.get_ram_cost_usage(in_size=(2, 16, 80, 80, 80), full=True) multi_res_images = [ torch.randn(2, 16, 80, 80, 80), torch.randn(2, 32, 40, 40, 40), torch.randn(2, 64, 20, 20, 20), torch.randn(2, 128, 10, 10, 10)] prediction = topology_search_space(image) for x in prediction: print(x.shape) # torch.Size([2, 16, 80, 80, 80]) # torch.Size([2, 32, 40, 40, 40]) # torch.Size([2, 64, 20, 20, 20]) # torch.Size([2, 128, 10, 10, 10]) Class method overview: - ``get_prob_a()``: convert learnable architecture weights to path activation probabilities. - ``get_ram_cost_usage()``: get estimated ram cost. - ``get_topology_entropy()``: get topology entropy loss in searching stage. - ``decode()``: get final binarized architecture code. - ``gen_mtx()``: generate variables needed for topology search. Predefined variables: - `tidx`: index used to convert path activation matrix T = (depth,depth) in transfer_mtx to path activation arch_code (1,3*depth-2), for depth = 4, tidx = [0, 1, 4, 5, 6, 9, 10, 11, 14, 15], A tidx (10 binary values) represents the path activation. - `transfer_mtx`: feasible path activation matrix (denoted as T) given a node activation pattern. It is used to convert path activation pattern (1, paths) to node activation (1, nodes) - `node_act_list`: all node activation [2^num_depths-1, depth]. For depth = 4, there are 15 node activation patterns, each of length 4. For example, [1,1,0,0] means nodes 0, 1 are activated (with input paths). - `all_connect`: All possible path activations. For depth = 4, all_connection has 1024 vectors of length 10 (10 paths). The return value will exclude path activation of all 0. z list[list]node2outnode2inrNrrKrLrMrNTrrrrrrrPrrTrUrVrWrrrrc  stj||||||||| | d g} tj} t| jdD]} | | d| j| d| d| d| q| _|\}}}t |_ fddtt j D_ |_t |_t jt jjf_tjD]6}tt jD],}j||fdkrjt||ft fddjjdjDj||f<qqwttjt jjdd j !_"ttjt jd d j !_#d d g_$dS) rrrXr7ci|] }tj||qSr)r node_act_listr;rrrr z+TopologySearch.__init__..csg|] }|jjjqSr)rFr)r;r) cell_interrrr=sz+TopologySearch.__init__..Ng{Gz?r log_alpha_a log_alpha_c)%r-r.rurrrrltidxgen_mtxrgasarrayrri node_act_dict transfer_mtx child_listzerosrrrrFrrrarrayrrer0 Parameterr/normal_rprrequires_grad_rr_arch_param_names)rrrrrrrTrUrWrrr_drrrrrrr3)rrrr.sT 6  &zTopologySearch.__init__depthc stj|d}td|d}g}|D]0}t||f}t|D]}||||dtj|dtjd|dtjf<q||qtd|ddd}i} |D]fdd|D} | | tt<qR| ||ddfS)a Generate elements needed in decoding and topology. - `transfer_mtx`: feasible path activation matrix (denoted as T) given a node activation pattern. It is used to convert path activation pattern (1, paths) to node activation (1, nodes) - `node_act_list`: all node activation [2^num_depths-1, depth]. For depth = 4, there are 15 node activation patterns, each of length 4. For example, [1,1,0,0] means nodes 0, 1 are activated (with input paths). - `all_connect`: All possible path activations. For depth = 4, all_connection has 1024 vectors of length 10 (10 paths). The return value will exclude path activation of all 0. rXrr7Ncs4g|]}t|ddkttkr|qSr9)rgsumastyperPr allr:rrrr=Ds4z*TopologySearch.gen_mtx..) rurr?rgrrrlrr ) rrrAZ all_connectmtxrmarrrZ arch_code_mtxrrrr"s 8 zTopologySearch.gen_mtxcsfddDS)Ncsg|] \}}|jvr|qSr)r rrrrr=Jsz4TopologySearch.weight_parameters..rrrrrrIsz TopologySearch.weight_parametersFrBcszt|jdddd}|r9t|j|jfddt |j D}t |}||fSd|fS)a Get final path and child model probabilities from architecture weights `log_alpha_a`. This is used in forward pass, getting training loss, and final decoding. Args: child: return child probability (used in decoding) Return: arch_code_prob_a: the path activation probability of size: `[number of blocks, number of paths in each block]`. For 12 blocks, 4 depths search space, the size is [12,10] probs_a: The probability of all child models (size 1023x10). Each child model is a path activation pattern (1D vector of length 10 for 10 paths). In total 1023 child models (2^10 -1) r7rcs:g|]}|dd|d|qS)r7r)prod)r;rZ_arch_code_prob_anormZpath_activationrrr=`s z-TopologySearch.get_prob_a..N) r/sigmoidrr unsqueezerrrprrrstack)rrBarch_code_prob_aprobs_arrr get_prob_aLs  zTopologySearch.get_prob_afullc Cs>|d}t||j d}g}t|jD]}|||j||d|qtj |tj |j ddt |j }|jdd\}}tj|jdd} |rW|}|d t|jjtj |j d} d } t|jD]+} tt|jD]!} | || | fd | | | f| | | f||j| 7} qtqk| d d d S)a Get estimated output tensor size to approximate RAM consumption. Args: in_size: input image shape (4D/5D, ``[BCHW[D]]``) at the highest resolution level. full: full ram cost usage with all probability of 1. rNrXrFrBrdimr7rori)rgr rbrrrlrrr/rfloat32rrPrrrsoftmaxrdetachfill_rrFrprrirr)rin_sizer batch_size image_sizesizesrrrZ cell_probrFusagerpath_idxrrrget_ram_cost_usagels0 &"   z!TopologySearch.get_ram_cost_usagec Cst|dr |j}|j}nddtt|jD}ddtt|jD}tt|jD]_}t|j t|j }}tt|j D]"}||j ||j||7<||j ||j||7<qA|dk t }|dk t }||jt||||jt||q+||_||_d}t|jdD]A} d} tt|jD]1} || || f} || d|| f} | | t| dd| td| d 7} q|| 7}q|S)z Get topology entropy loss at searching stage. Args: probs: path activation probabilities rcSg|]}gqSrrr:rrrr=z7TopologySearch.get_topology_entropy..cSr/rrr:rrrr=r0r7rh㈵>)hasattrrrrrirrrgrrrrrrPrrrlrrr/log)rprobsrr child_idxZ_node_inZ _node_outrentrZblk_entnode_idxZ_node_pZ _out_probsrrrget_topology_entropys2  "4 z#TopologySearch.get_topology_entropyc sjdd\}}jt|dj}ttj ddj}|j}t dt jj ddt jj df}fddtt jD}t t jt jf}tt jD]L}t jt} tt j|D]} | j| j|| 7<q~| dkt} jt| D]} |t| jt} d||| f<qqjt |dd d |dddt jf<tdj D]A} |t t || d d t jdf|d| dt jd| t jd| t jd| dt jf<qd |dj dt jdj t jdf<t|}t|ddddd \}}}d \}}t j t jf}t j djf} ||}|dkren6|dt j}j|||d d f<tt jD]}||j|f|||f7<q|d8}q[tt jD]}||j|f|d|f7<q|dkt}||||fS) a Decode network log_alpha_a/log_alpha_c using dijkstra shortest path algorithm. `[node_a, arch_code_a, arch_code_c, arch_code_a_max]` is decoded when using ``self.decode()``. For example, for a ``num_depths=4``, ``num_blocks=12`` search space: - ``node_a`` is a 4x13 binary matrix representing if a feature node is activated (13 because of multi-resolution inputs). - ``arch_code_a`` is a 12x10 (10 paths) binary matrix representing if a path is activated. - ``arch_code_c`` is a 12x10x5 (5 operations) binary matrix representing if a cell operation is used. Return: arch_code with maximum probability Tr rr7crr)rrrrrrrrz)TopologySearch.decode..rr1gMbP?)csgraphdirectedindicesmin_onlyreturn_predecessors)rrN)rrr/argmaxrrnumpyrr%rrgrrirrrrrPrrrflattenrr3tiler rr)rr4rZarch_code_a_maxrmZamtxZ path2childZsub_amtxr5Z _node_actr-rZconnect_child_idxrgraph dist_matrix predecessorssourcesindexZa_idxrrrrrrdecodes` .".,&&4   & & zTopologySearch.decodec Cs|jdd\}}|}t|jD]J}dg|j}t|j|jD]3\}}|rVt j |j ||fdd} ||j ||j t||f||j|| d|||f7<q#|}q|S)z Prediction based on dynamic arch_code. Args: x: a list of `num_depths` input tensors as a multi-resolution input. tensor is of shape `BCHW[D]` where `C` must match `self.filter_nums`. Fr rorr!)r)rrrrrqrrrr?rr%rrrrr) rrrrrrrrr_wrrrr s   " zTopologySearch.forward)rrrrrrPrrPrTrPrUrVrWrVrrrr)rrP)F)rBr)rr)r$r%r&r'__annotations__rur.rrrr.r8rGr r6rrr3rrs, -  >'  !&Lr)1 __future__rrrtypingrr?rgr/torch.nnr0torch.nn.functional functionalrZ!monai.networks.blocks.dints_blockrrrrmonai.networks.layers.factoriesrmonai.networks.layers.utilsr r monai.utilsr r r<r__all__jit interfaceModulerr)r,r?IdentityrCrJrYr]rarcrurrrrrrrrsB      "{Ju8