U Ph@sddlmZddlZddlZddlmZddlZddlZddl m Z ddl m m Z ddlmZmZmZmZddlmZddlmZmZddlmZedd d \ZZed d d \ZZd dddgZejjGdddej j Z!ejjGdddej j Z"Gddde"Z#ddZ$Gddde j%Z&GdddeZ'GdddeZ(Gdd d eZ)Gd!d"d"eZ*Gd#d$d$e j Z+Gd%d&d&e!Z,Gd'd d e j Z-Gd(dde j Z.Gd)dde.Z/Gd*dde.Z0dS)+) annotationsN)Optional)ActiConvNormBlockFactorizedIncreaseBlockFactorizedReduceBlockP3DActiConvNormBlock)Conv) get_act_layerget_norm_layer)optional_importz scipy.sparse csr_matrixnamezscipy.sparse.csgraphdijkstraDiNTSTopologyConstructionTopologyInstanceTopologySearchc@s"eZdZdZddddddZdS) CellInterfacez"interface for torchscriptable Cell torch.TensorOptional[torch.Tensor]xweightreturncCsdSNselfrrrrN/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/dints.pyforward,szCellInterface.forwardN__name__ __module__ __qualname____doc__r rrrrr(src@s eZdZdZdddddZdS) StemInterfacez"interface for torchscriptable StemrrrcCsdSrrrrrrrr 4szStemInterface.forwardNr!rrrrr&0sr&cs0eZdZdZfddZdddddZZS)StemTSz wrapper for torchscriptable Stemcsttjj||_dSr)super__init__torchnn Sequentialmod)rr/ __class__rrr+;s zStemTS.__init__rr'cCs ||Sr)r/r(rrrr ?szStemTS.forwardr"r#r$r%r+r __classcell__rrr0rr)8s r)cCs>||krdgdggSt|d|}dd|Ddd|DS)z>use depth first search to find all path activation combinationrcSsg|]}dg|qSrr.0_rrr Hsz_dfs..cSsg|]}dg|qS)r4rr6rrrr9Hs)_dfs)nodepathschildrrrr:Cs r:cseZdZfddZZS)_IdentityWithRAMCostcstj||d|_dSNrr*r+ram_cost)rargskwargsr0rrr+Msz_IdentityWithRAMCost.__init__r"r#r$r+r3rrr0rr>Ksr>c sBeZdZdZdddddiffdddddddd fd d 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 RELUINSTANCEaffineTint tuple | str) in_channel out_channel kernel_sizepadding spatial_dimsact_name norm_namec s.t|||||||d||d|_dS)Nr4r@)rrLrMrNrOrPrQrRr0rrr+^s z&_ActiConvNormBlockWithRAMCost.__init__)r"r#r$r%r+r3rrr0rrERs  rEc s>eZdZdddddiffddddddddfd d ZZS) _P3DActiConvNormBlockWithRAMCostrrGrHrITrJrK)rLrMrNrOp3dmoderQrRc s.t|||||||dd|||_dSNrSr@)rrLrMrNrOrUrQrRr0rrr+ns z)_P3DActiConvNormBlockWithRAMCost.__init__rDrrr0rrTls rTcs:eZdZdddddiffddddddfd d ZZS) #_FactorizedIncreaseBlockWithRAMCostrFrGrHrITrJrKrLrMrPrQrRcs*t|||||d||d|_dSrVr@rrLrMrPrQrRr0rrr+sz,_FactorizedIncreaseBlockWithRAMCost.__init__rDrrr0rrW~s rWcs:eZdZdddddiffddddddfd d ZZS) !_FactorizedReduceBlockWithRAMCostrFrGrHrITrJrKrXcs0t|||||||d|jd|_dS)NrSrF)r*r+ _spatial_dimsrArYr0rrr+sz*_FactorizedReduceBlockWithRAMCost.__init__rDrrr0rrZs rZcs<eZdZdZd dddfdd Zddd d d d ZZS)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. NrJdict)copscs^t|dkr tt|}t|_t||D]$\}}|dkr4|j |||q4dSr?) r*r+nponeslenr- ModuleListr_zipappend)rr^r_ arch_code_cZarch_cop_namer0rrr+s  zMixedOp.__init__rrrrcCsVd}|dk r||}t|jD]0\}}|dkr<|||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 enumerater_)rrroutidx_oprrrr s  *zMixedOp.forward)N)Nr2rrr0rr\s r\cseZdZdZdZdddddZddddd dd dd dd Zeee e d Z ddddddiffdddddddfdd Z ddddddZ 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 rFcCstSrr>_crrrz Cell.cCst||ddddS)NrFr4rSrOrPrEr^rrrrsrt skip_connectZconv_3x3cCstSrrprqrrrrsrtcCst||ddddS)NrFr4rurvrwrrrrsrtcCst||ddddS)NrFr4rrOrUrTrwrrrrsrtcCst||ddddS)NrFr4rzr{rwrrrrsrtcCst||ddddS)NrFr4rSrzr{rwrrrrsrtryZ conv_3x3x3Z conv_3x3x1Z conv_3x1x3Z conv_1x3x3)updownidentityalign_channelsNrGrHrITrJrK)c_prevr^raterPrQrRc sRt|_|_|_|dkrFjd||jjjd_nf|dkrpjd||jjjd_n<||krjd_n$jd||ddjjjd_d d fd d d _d d fdd fdd fdd fdd d_i_ jdkrj_ n(jdkr,j_ nt djdt |j |_ dS)Nr~)rPrQrRr4r}rrrcSstSrrprqrrrrsrtzCell.__init__..c st||dddjjdS)NrFr4rSrOrPrQrRrE _act_name _norm_namerwrrrrssrxcSstSrrprqrrrrsrtc st||dddjjdS)NrFr4rrrwrrrrssc st||dddjjdS)NrFr4rrOrUrQrRrTrrrwrrrrssc st||dddjjdS)NrFr4rrrwrrrrs sc st||dddjjdS)NrFr4rSrrrwrrrrs#sr|rSrFSpatial dimensions  is not supported.) r*r+r[rrConnOPS preprocessOPS2DOPS3DZOPSNotImplementedErrorr\op)rrr^rrfrPrQrRr0rrr+s^             z Cell.__init__rrrcCs||}|||}|S)zi Args: x: input tensor weight: weights for different operations. )rrrrrrr 2s  z Cell.forward)r"r#r$r% DIRECTIONSrrrWrZr>rErr+r r3rrr0rros,     BrocsZeZdZdZddddifdddfddd d dd d fd d ZddZddddZZS)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. rGrHrITrFNrJrKbool) in_channels num_classesrQrRrPuse_downsamplec st||_|j|_|j|_|j|_|dkr@td|d||_|dkrht |jd|jf|_ n||_ t t j |f} t |_t |_t t|jd|jd|||d| |jd|ddddddd|_|jd krd 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| <qtt jdd | | dd | ||j| d ddddddt|||j| d|jt| <tt|d| |j| |jt| ddd ddddddt|||jt| dddt jd | dk| dd |jt| <qdS)N)rSrFrrr4r)rQrRrPT)r out_channelsrNstriderOgroupsbiasdilationrF trilinearbilinearrS) scale_factormode align_cornersF)rrPchannelsr )r*r+ dints_space filter_nums num_blocks num_depthsrr[r,ranode_arCONVr- ModuleDict stem_downstem_upr.r stem_finalsranger)Upsampler r strmax) rrrrrQrRrPrr conv_typerres_idxr0rrr+cs          zDiNTS.__init__cCsdd|DS)NcSsg|] \}}|qSrrr7rparamrrrr9sz+DiNTS.weight_parameters..named_parametersrrrrweight_parametersszDiNTS.weight_parametersr)rc Csg}t|jD]F}|jt|}||}|jd|rD||q|t|q| |}|j d}d}t d} t|jdddD]N} |j t| } |r| || | } q|j|d| rd}| || } q| | } | S)zd Prediction based on dynamic arch_code. Args: x: input tensor. rr4FrT)rrrrr rrer, zeros_likerremptyrr) rrinputsdZ_mod_wx_outoutputsblk_idxstart_temprZ_mod_up predictionrrrr s(      z DiNTS.forward)r"r#r$r%r+rr r3rrr0rr=s* c s\eZdZdZddeddddddd ifd d f d d d d d ddddd fdd 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?rFrGrHrITcpu list | NonefloatrJrKrr) arch_code channel_mulrrrPrQrRrdevicec  sttddt|dD} fdd| D|_||_||_tdt j d|t | ||_ ||_ ||_| |_| |_d|_|j dkrt|j|_n|j dkrt|j|_gg} } ttj|jdD]*}| |dtjd|dtjqd ddg|jdd }t|jD]}| |||gq$| dd } | |_||_| |_|dkrt|jt|jf|j}t|jt|j|jf|j}n>t|d|j}t !t|dtj"|j|j}||_#||_$t%&|_'t|jD]}tt|jD]z}|j#||fdkr||j|j|t | |j|j|t | |j||j$||f|j |j |j|j't(||f<qqdS) NcSsg|]}dd|qS) rSr)r7_irrrr9;sz1TopologyConstruction.__init__..r4csg|]}t|qSr)rJ)r7Zn_featrrrr9<szC{} - Length of input patch is recommended to be a multiple of {:d}.rSrrFr))r*r+tuplerrrrprintformatdatetimenowrJr[rrrr num_cell_opsrbrrrorreextend arch_code2in arch_code2ops arch_code2outr,rarj from_numpyFone_hotint64 arch_code_arfr-r cell_treer)rrrcellrrrPrQrRrrZn_featsrrirmrrfrrr0rrr+,sf     (  $(  zTopologyConstruction.__init__cCsdS)zOThis function to be implemented by the architecture instances or search spaces.Nrr(rrrr vszTopologyConstruction.forwardr"r#r$r%ror+r r3rrr0rrs' $Jc sbeZdZdZddeddddddd ifd d f d d d d d d dddfdd ZdddddZZS)rz` Instance of the final searched architecture. Only used in re-training/inference stage. NrrrFrGrHrITrrrJrKrr)rrrrPrQrRrrc s6|dkrtdtj||||||||| | d dS)Q Initialize DiNTS topology search space of neural architectures. Nz*arch_code not provided when not searching. rrrrrrPrQrRrr)warningswarnr*r+) rrrrrrrPrQrRrrr0rrr+s zTopologyInstance.__init__zlist[torch.Tensor]r'c Cs|}t|jD]}tjd|dj|djdg|j}t|j|j D]R\}}|rF|j t ||f}|j ||j |dd}||j||||j|<qF|}q|S)z4 Args: x: input tensor. rirdtyperNrh)rrr,tensorrrrrkrdatarrr rr) rrrrrr activationr/_outrrrr s$zTopologyInstance.forwardrrrr0rr{s " c seZdZUdZded<ded<deddddd d d d ifd d f dddddddddd fdd ZddddZddZd)ddddZ d*dd d!d"Z d#d$Z d%d&Z d'd(Z ZS)+ra 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]node2outnode2inrNrrFrGrHrITrrrrJrKrr) rrrrrPrQrRrrc  stj||||||||| | d g} tj} t| jdD]8} | | d| j| d| d| d| q<| _|\}}}t |_ fddtt j D_ |_t |_t jt jjf_tjD]ptt jD]\jfdkrt fddjtfjjdjDjf<qqttjt jjdd j !_"ttjt jd d j !_#d d g_$dS) rrrSr4csi|]}tj||qSr)r node_act_listr7rrrr s z+TopologySearch.__init__..cs(g|] }|jjtfjjqSr)rArrr)r7rrrrrrr9sz+TopologySearch.__init__..Ng{Gz?r log_alpha_a log_alpha_c)%r*r+rorrrretidxgen_mtxr`asarrayrrb node_act_dict transfer_mtx child_listzerosrrrrArarrayrrrr_r- Parameterr,normal_rjrrequires_grad_rr_arch_param_names)rrrrrrrPrQrRrrr_drrrrr0rrr+sV 6  &zTopologySearch.__init__)depthc stj|d}td|d}g}|D]`}t||f}t|D]:}||||dtj|dtjd|dtjf<q>||q$td|ddd}i} |D](fdd|D} | | tt<q| ||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. rSrr4Ncs4g|],}t|ddkttkr|qSr5)r`sumastyperJrallr6rrrr9Fs$z*TopologySearch.gen_mtx..) rorr:r`rrrerr) rrr<Z all_connectmtxrmarrrZ arch_code_mtxrrrr$s 8 zTopologySearch.gen_mtxcsfddDS)Ncsg|]\}}|jkr|qSr)rrrrrr9Ls z4TopologySearch.weight_parameters..rrrrrrKsz TopologySearch.weight_parametersFr=cszt|jdddd}|rrt|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) r4rcs:g|]2}|dd|d|qS)r4r)prod)r7rZ_arch_code_prob_anormZpath_activationrrr9bs z-TopologySearch.get_prob_a..N) r,sigmoidrr unsqueezerrrjrrrstack)rr=arch_code_prob_aprobs_arr r get_prob_aNs  zTopologySearch.get_prob_a)fullc Cs>|d}t||j d}g}t|jD](}|||j||d|q,tj |tj |j ddt |j }|jdd\}}tj|jdd} |r|}|d t|jjtj |j d} d } t|jD]V} tt|jD]B} | || | fd | | | f| | | f||j| 7} qq| 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. rNrSrFrrdimr4riri)r`rr[rrrerrr,rfloat32rrJrrrsoftmaxrdetachfill_rrArjrrbrr)rin_sizer batch_size image_sizesizesrrrZ cell_probrAusagerpath_idxrrrget_ram_cost_usagens. &"   z!TopologySearch.get_ram_cost_usagec Cst|dr|j}|j}n ddtt|jD}ddtt|jD}tt|jD]}t|j t|j }}tt|j D]D}||j ||j||7<||j ||j||7<q|dk t }|dk t }||jt||||jt||qX||_||_d}t|jdD]} d} tt|jD]d} || || f} || d|| f} | | t| dd| td| d 7} qL|| 7}q6|S)z Get topology entropy loss at searching stage. Args: probs: path activation probabilities rcSsg|]}gqSrrr6rrrr9sz7TopologySearch.get_topology_entropy..cSsg|]}gqSrrr6rrrr9sr4rh㈵>)hasattrrrrrbrrr`rrrrrrJrrrerrr,log)rprobsrr child_idxZ_node_inZ _node_outrentrZblk_entnode_idxZ_node_pZ _out_probsrrrget_topology_entropys2   "6 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]}t jt} tt j|D]$} | j| j|| 7<q| dkt} jt| D].} |t| jt} d||| f<q>qt |dd d |dddt jf<tdj D]} |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}||}|dkrΐq:|dt j}j|||d d f<tt jD](}||j|f|||f7<q|d8}qtt jD](}||j|f|d|f7<qH|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 Trrr4csi|]}tj||qSr)rrrrrrrs z)TopologySearch.decode..rr gMbP?)csgraphdirectedindicesmin_onlyreturn_predecessors)rrN)rrr,argmaxrrnumpyrrrr`rrbrrrrrJrrrflattenrr"tiler rr)rr#rZarch_code_a_maxrfZamtxZ path2childZsub_amtxr$Z _node_actrrZconnect_child_idxrgraph dist_matrix predecessorssourcesindexZa_idxrrrrrrdecodesf ."., 4  & &zTopologySearch.decodec Cs|jdd\}}|}t|jD]}dg|j}t|j|jD]f\}}|rFt j |j ||fdd} ||j ||j t||f||j|| d|||f7<qF|}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`. Frrirr)r)rrrrrkrrrr.rrrrrrr) rrrrrrrrr_wrrrr s   " zTopologySearch.forward)F)F)r"r#r$r%__annotations__ror+rrrrr'r6r r3rrr0rrs* - $@' !&L)1 __future__rrrtypingrr.r`r,torch.nnr-torch.nn.functional functionalrZ!monai.networks.blocks.dints_blockrrrrmonai.networks.layers.factoriesrmonai.networks.layers.utilsr r monai.utilsr r r8r__all__jit interfaceModulerr&r)r:Identityr>rErTrWrZr\rorrrrrrrr s@       "{Ju8