o  i@s(ddlmZddlZddlZddlmZddlZddlmZddl mm Z ddl m Z ddlmZddlmZmZddlmZeeZdd d ZGdddejZGdddejZGdddejZGdddejZGdddejZGdddejZ GdddejZ!GdddeZ"dS)) annotationsN)Sequence) Convolution)SpatialAttentionBlock) AEKLResBlock AutoencoderKL)convert_to_tensorsave_memboolreturnNonecCstjr |r tjdSN)torchcuda is_available empty_cache)r rz/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py_empty_cuda_caches rcs@eZdZdZ     ddfdd ZdddZddZZS)MaisiGroupNorm3DaY Custom 3D Group Normalization with optional print_info output. Args: num_groups: Number of groups for the group norm. num_channels: Number of channels for the group norm. eps: Epsilon value for numerical stability. affine: Whether to use learnable affine parameters, default to `True`. norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. h㈵>TF num_groupsint num_channelsepsfloataffiner norm_float16 print_infor cs(t||||||_||_||_dSr )super__init__rrr )selfrrrrrrr  __class__rrr 3s  zMaisiGroupNorm3D.__init__input torch.Tensorr c Cs|jr td|t|jdkrtd|j\}}}}}|||j||j|||}g}t |dD]H}|dd||ddfj t j d} | j gddd } | jgdd dd |j} |jrv|| | | j t jdq7|| | | q7~t|jt|d d krt j|ddn||}||||||}|jr||jd|ddd|jd|ddd|jrtd||S)Nz"MaisiGroupNorm3D with input size: zExpected a 5D tensor.)dtype)r&T)keepdimF)unbiasedr,rdimz#MaisiGroupNorm3D with output size: )rloggerinfosizelenshape ValueErrorviewrrangetorfloat32meanvaradd_rsqrt_rappendfloat16rr maxcat _cat_inputsrmul_weightbias) r!r$Zparam_nZparam_cZparam_dZparam_hZparam_winputsiarrayr;stdrrrforwardBs.$  ,0zMaisiGroupNorm3D.forwardcCs|djj}|dkr|djdddn|d}d|d<t|jtt|dD]5}tj |||d fdd}d||d<t|jt |j ratd|dd t|dd q,|dkrm|jdddS|S) NrrcpuT non_blockingr'r/z"MaisiGroupNorm3D concat progress: /.)devicetypecloner9rr r8r4rrBrLgccollectrr1r2)r!rG input_typer$krrrrCds *   $zMaisiGroupNorm3D._cat_inputs)rTFFT)rrrrrrrr rr rr r r )r$r%r r%)__name__ __module__ __qualname____doc__r rKrC __classcell__rrr"rr%s "rcs`eZdZdZ             d5d6fd(d) Zd7d.d/Zd8d1d2Zd9d3d4ZZS):MaisiConvolutiona Convolutional layer with optional print_info output and custom splitting mechanism. Args: spatial_dims: Number of spatial dimensions (1D, 2D, 3D). in_channels: Number of input channels. out_channels: Number of output channels. num_splits: Number of splits for the input tensor. dim_split: Dimension of splitting for the input tensor. print_info: Whether to print information. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. Additional arguments for the convolution operation. https://docs.monai.io/en/stable/networks.html#convolution Tr'r*NDAPRELUINSTANCENF spatial_dimsr in_channels out_channels num_splits dim_splitrr r stridesSequence[int] | int kernel_size adn_orderingstracttuple | str | Nonenormdropouttuple | str | float | None dropout_dimdilationgroupsrF conv_only is_transposedpaddingSequence[int] | int | Noneoutput_paddingr r csttdid|d|d|d|d| d| d| d| d | d |d |d |d |d|d|d|d||_||_t|trK||jn||_||_||_ ||_ dS)NrarbrcrfrhrirkrmrnrprqrrrFrsrtrurwr) rr rconvre isinstanceliststriderdrr )r!rarbrcrdrerr rfrhrirkrmrnrprqrrrFrsrtrurwr"rrr sR       zMaisiConvolution.__init__xr% split_sizelist[torch.Tensor]c Csdg|g|jd}||jd|}tdgd}g}t|jD]*}t|||||d|||jdkr;|n|||jd<||t|q#|jrptt|D]} t d| ddt|d|| qW|S)Nrr'r)r&zSplit rOz size: ) rdr3reslicer8r?tuplerr4r1r2) r!r|r}ruoverlapsZ last_paddingslicessplitsrHjrrr _split_tensors .zMaisiConvolution._split_tensoroutputsc Cstdgd}t|jD]!}|dkrtd|nt|||||jd<||t|||<q |jrOt|jD]}td|ddt|d|| q6t |d dkret j ||jdd }|S|d jd d d }t d|d<t|jtt|dD];}t j |||df|jdd }t d||d<t|jt|jrtd |ddt|ddq|jdd d }|S)Nr&rr)Output r'rOz size after: r.r/rLTrMz"MaisiConvolution concat progress: rPr)rr8rdrerrr1r2r4r3rArrBrSr9Tensorrr rLrTrU)r!rr}rurrHr|rWrrr_concatenate_tensorss.*. $ $z%MaisiConvolution._concatenate_tensorsc sjr tdjjdkr|}|S|jd}|j}d}|jdkr8|jdj}jrCtd||||}~t j fdd|D}jr{t t |D]}td |dd t |d ||qb|}|} jdkrjdnd} |d| d|d| ddkr|d9}| d9} n|d| d|d| ddkr|d}| d} ||| }~t j |S) NzNumber of splits: r'r)r*rzPadding size: csg|]}|qSr)rx).0splitr!rr sz,MaisiConvolution.forward..rrOz size before: )rr1r2rdrxr3rer{rrr r8r4r) r!r|lr}rurrrZsplit_size_outZ padding_sZ non_dim_splitrrrrKs@    .( ( zMaisiConvolution.forward)Tr'r*r^r_r`Nr'r'r'TFFNN),rarrbrrcrrdrrerrr r r rfrgrhrgrirjrkrlrmrlrnrorprrqrgrrrrFr rsr rtr rurvrwrvr r )r|r%r}rrurr r~)rr~r}rrurr r%r|r%r r%) rXrYrZr[r rrrKr\rrr"rr]vs* 3 r]cs0eZdZdZ ddfdd ZdddZZS) MaisiUpsamplea Convolution-based upsampling layer. Args: spatial_dims: Number of spatial dimensions (1D, 2D, 3D). in_channels: Number of input channels to the layer. use_convtranspose: If True, use ConvTranspose to upsample feature maps in decoder. num_splits: Number of splits for the input tensor. dim_split: Dimension of splitting for the input tensor. print_info: Whether to print information. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. Trarrbuse_convtransposer rdrerr r r csDtt||||r dndddd|||||d |_||_||_dS)Nr)r'r*T) rarbrcrfrhrursrtrdrerr )rr r]rxrr )r!rarbrrdrerr r"rrr (s"   zMaisiUpsample.__init__r|r%cCsV|jr||}t|}|Stj|ddd}t|j||}t|jt|}|S)Ng@ trilinear) scale_factormode)rrxrF interpolaterr )r!r|x_tensor out_tensorrrrrKDs    zMaisiUpsample.forwardT)rarrbrrr rdrrerrr r r r r rrXrYrZr[r rKr\rrr"rrs rcs0eZdZdZ ddfd d ZdddZZS)MaisiDownsamplea Convolution-based downsampling layer. Args: spatial_dims: Number of spatial dimensions (1D, 2D, 3D). in_channels: Number of input channels. num_splits: Number of splits for the input tensor. dim_split: Dimension of splitting for the input tensor. print_info: Whether to print information. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. Trarrbrdrerr r r r c s8td||_t|||dddd||||d |_dS)N)rr'r)r*rT rarbrcrfrhrursrdrerr )rr padr]rx)r!rarbrdrerr r"rrr `s  zMaisiDownsample.__init__r|r%cCs"tj||jddd}||}|S)Nconstantg)rvalue)rrrx)r!r|rrrrKys zMaisiDownsample.forwardr)rarrbrrdrrerrr r r r r rrrrr"rrSs rcs4eZdZdZ   ddfdd ZdddZZS) MaisiResBlockaJ Residual block consisting of a cascade of 2 convolutions + activation + normalisation block, and a residual connection between input and output. Args: spatial_dims: Number of spatial dimensions (1D, 2D, 3D). in_channels: Input channels to the layer. norm_num_groups: Number of groups for the group norm layer. norm_eps: Epsilon for the normalization. out_channels: Number of output channels. num_splits: Number of splits for the input tensor. dim_split: Dimension of splitting for the input tensor. norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. FTrarrbnorm_num_groupsnorm_epsrrcrdrerr rr r r c st||_|dur|n||_| |_t|||d|| | d|_t||j|jdddd||| | d |_t|||d|| | d|_ t||j|jdddd||| | d |_ |j|jkrjt||j|jdddd||| | d |_ dSt |_ dS)NTrrrrrrr r'r*rr)rr rbrcr rnorm1r]conv1norm2conv2nnIdentity nin_shortcut) r!rarbrrrcrdrerrr r"rrr s     zMaisiResBlock.__init__r|r%cCs||}t|jt|}t|j||}t|j||}t|jt|}t|j||}t|j|j|j krL| |}t|j||}t |}|Sr ) rrr rsilurrrrbrcrr)r!r|houtrrrrrKs$               zMaisiResBlock.forward)FFT)rarrbrrrrrrcrrdrrerrr rr r r r r rrrrr"rrsQrcs<eZdZdZ       d"d#fdd Zd$d d!ZZS)% MaisiEncodera Convolutional cascade that downsamples the image into a spatial latent space. Args: spatial_dims: Number of spatial dimensions (1D, 2D, 3D). in_channels: Number of input channels. num_channels: Sequence of block output channels. out_channels: Number of channels in the bottom layer (latent space) of the autoencoder. num_res_blocks: Number of residual blocks (see AEKLResBlock) per level. norm_num_groups: Number of groups for the group norm layers. norm_eps: Epsilon for the normalization. attention_levels: Indicate which level from num_channels contain an attention block. with_nonlocal_attn: If True, use non-local attention block. include_fc: whether to include the final linear layer in the attention block. Default to False. use_combined_linear: whether to use a single linear layer for qkv projection in the attention block, default to False. use_flash_attention: If True, use flash attention for a memory efficient attention mechanism. num_splits: Number of splits for the input tensor. dim_split: Dimension of splitting for the input tensor. norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. FTrarrbr Sequence[int]rcnum_res_blocksrrrattention_levelsSequence[bool]rdrerr rr with_nonlocal_attn include_fcuse_combined_linearuse_flash_attentionr r cstt|t|krtdt|t|krtd| |_g}|t|||ddddd| | | | d |d}tt|D]M}|}||}|t|dk}t||D]'}|t|||||| | | | | d |}||r}|t |||||||d qV|s|t ||| | | | d q@|r|t ||d |||d d |t ||d |||||d |t ||d |||d d |t ||d |d| | | d |t||d |dddd| | | | d t ||_dS)Nz9attention_levels and num_channels must have the same sizez7num_res_blocks and num_channels must have the same sizerr'r*Tr rarbrrrcrdrerrr rarrrrrr)rarbrdrerr rarbrrrcr)rr r4r6r r?r]r8rrrrrr ModuleListblocks)r!rarbrrcrrrrrdrerrr rrrrroutput_channelrH input_channelis_final_block_r"rrr s       zMaisiEncoder.__init__r|r%cC"|jD] }||}t|jq|Sr rrr r!r|blockrrrrK  zMaisiEncoder.forward)FFTTFFF)$rarrbrrrrcrrrrrrrrrrdrrerrr rr r r rr rr rr rr r r rrrrr"rrs#rcs>eZdZdZ        d#d$fdd Zd%d!d"ZZS)& MaisiDecodera Convolutional cascade upsampling from a spatial latent space into an image space. Args: spatial_dims: Number of spatial dimensions (1D, 2D, 3D). num_channels: Sequence of block output channels. in_channels: Number of channels in the bottom layer (latent space) of the autoencoder. out_channels: Number of output channels. num_res_blocks: Number of residual blocks (see AEKLResBlock) per level. norm_num_groups: Number of groups for the group norm layers. norm_eps: Epsilon for the normalization. attention_levels: Indicate which level from num_channels contain an attention block. with_nonlocal_attn: If True, use non-local attention block. include_fc: whether to include the final linear layer in the attention block. Default to False. use_combined_linear: whether to use a single linear layer for qkv projection in the attention block, default to False. use_flash_attention: If True, use flash attention for a memory efficient attention mechanism. use_convtranspose: If True, use ConvTranspose to upsample feature maps in decoder. num_splits: Number of splits for the input tensor. dim_split: Dimension of splitting for the input tensor. norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. FTrarrrrbrcrrrrrrrdrerr rr rrrrrr r cst| |_| |_tt|}g}|t|||ddddd| | | | d |rY|t||d|||dd|t ||d|||||d|t||d|||ddtt|}tt|}|d}t t |D]N}|}||}|t |dk}t ||D]'}|t |||||| | | | | d |}||r|t |||||||dq|s|t |||| | | | d qo|t|||d| | | d |t|||dddd| | | | d t||_dS) Nrr'r*Trrrr)rarbrrdrerr r)rr rr rzreversedr?r]rrr8r4rrrrrr)r!rarrbrcrrrrrdrerrr rrrrrreversed_block_out_channelsrreversed_attention_levelsreversed_num_res_blocks block_out_chrH block_in_chrrr"rrr s        zMaisiDecoder.__init__r|r%cCrr rrrrrrKnrzMaisiDecoder.forward)FFTTFFFF)&rarrrrbrrcrrrrrrrrrrdrrerrr rr r r rr rr rr rr rr r r rrrrr"rrs$rcsBeZdZdZ               d'd(fd%d& ZZS))AutoencoderKlMaisiar AutoencoderKL with custom MaisiEncoder and MaisiDecoder. Args: spatial_dims: Number of spatial dimensions (1D, 2D, 3D). in_channels: Number of input channels. out_channels: Number of output channels. num_res_blocks: Number of residual blocks per level. num_channels: Sequence of block output channels. attention_levels: Indicate which level from num_channels contain an attention block. latent_channels: Number of channels in the latent space. norm_num_groups: Number of groups for the group norm layers. norm_eps: Epsilon for the normalization. with_encoder_nonlocal_attn: If True, use non-local attention block in the encoder. with_decoder_nonlocal_attn: If True, use non-local attention block in the decoder. include_fc: whether to include the final linear layer. Default to False. use_combined_linear: whether to use a single linear layer for qkv projection, default to False. use_flash_attention: If True, use flash attention for a memory efficient attention mechanism. use_checkpointing: If True, use activation checkpointing. use_convtranspose: If True, use ConvTranspose to upsample feature maps in decoder. num_splits: Number of splits for the input tensor. dim_split: Dimension of splitting for the input tensor. norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. r* ư>FrTrarrbrcrrrrrlatent_channelsrrrwith_encoder_nonlocal_attnr with_decoder_nonlocal_attnrrruse_checkpointingrrdrerrr r r cst||||||||| | | ||| | |tdid|d|d|d|d|d|d| d|d | d | d | d |d |d|d|d|d||_tdid|d|d|d|d|d|d| d|d | d | d | d |d|d |d|d|d|d||_dS)Nrarbrrcrrrrrrrrrdrerrr rr)rr rencoderrdecoder)r!rarbrcrrrrrrrrrrrrrrdrerrr r"rrr s           zAutoencoderKlMaisi.__init__)r*rrFFFFFFFrrFFT),rarrbrrcrrrrrrrrrrrrrrr rr rr rr rr rr rr rdrrerrr rr r r r r )rXrYrZr[r r\rrr"rrus$#r)r r r r )# __future__rrTloggingcollections.abcrrtorch.nnrtorch.nn.functional functionalrmonai.networks.blocksrZ&monai.networks.blocks.spatialattentionrZ!monai.networks.nets.autoencoderklrrmonai.utils.type_conversionr getLoggerrXr1r GroupNormrModuler]rrrrrrrrrrs0        Q%9,}>=