o ( i@sddlmZddlmZddlZddlmZddlmmZ ddl m Z ddl m Z e ddd\ZZd d gZGd d d ejZGd d d ejZdS) ) annotations)castN) Convolution)optional_importeinops rearrange)name FeedForwardCABlockcs,eZdZdZdfd d ZdddZZS)r aGated-DConv Feed-Forward Network (GDFN) that controls feature flow using gating mechanism. Uses depth-wise convolutions for local context mixing and GELU-activated gating for refined feature selection. Args: spatial_dims: Number of spatial dimensions (2D or 3D) dim: Number of input channels ffn_expansion_factor: Factor to expand hidden features dimension bias: Whether to use bias in convolution layers spatial_dimsintdimffn_expansion_factorfloatbiasboolc srtt||}t|||dd|dd|_t||d|dddd|d|dd |_t|||d|dd|_dS)NTr in_channels out_channels kernel_sizer conv_only r rrrstridespaddinggroupsrr)super__init__r r project_indwconv project_out)selfr r rrZhidden_features __class___/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/cablock.pyr&s:     zFeedForward.__init__x torch.TensorreturncCs>||}||jddd\}}ttj|t||S)Nrrr ) r r!chunkrtorchTensorr"Fgelu)r#r(x1x2r&r&r'forwardHs zFeedForward.forward)r r r r rrrrr(r)r*r))__name__ __module__ __qualname____doc__rr3 __classcell__r&r&r$r'r s "csFeZdZdZddfd d Zd d Zd dZddZdddZZ S)r a5Multi-DConv Head Transposed Self-Attention (MDTA): Differs from standard self-attention by operating on feature channels instead of spatial dimensions. Incorporates depth-wise convolutions for local mixing before attention, achieving linear complexity vs quadratic in vanilla attention. Based on SW Zamir, et al., 2022 Args: spatial_dims: Number of spatial dimensions (2D or 3D) dim: Number of input channels num_heads: Number of attention heads bias: Whether to use bias in convolution layers flash_attention: Whether to use flash attention optimization. Defaults to False. Raises: ValueError: If flash attention is not available in current PyTorch version ValueError: If spatial_dims is greater than 3 Fr r num_headsrrflash_attentionc st|rttdstd|dkrtd|||_||_tt |dd|_ ||_ t |||dd|dd|_t ||d|dddd|d|dd |_t |||d|dd|_||_dS) Nscaled_dot_product_attentionzFlash attention not availablerz6Only 2D and 3D inputs are supported. Got spatial_dims=rTrr)rrhasattrr/ ValueErrorr r:nn Parameterr-ones temperaturer;rqkv qkv_dwconvr"_get_attention_fn _attention_fn)r#r r r:rr;r$r&r'r`s6  zCABlock.__init__cCs|jr|jS|jS)N)r;_flash_attention_normal_attention)r#r&r&r'rEszCABlock._get_attention_fncCs(t|j}tj||||ddd}|S)zBFlash attention implementation using scaled dot-product attention.gF)scale dropout_p is_causal)rrBmeanr/r<)r#qkvrIoutr&r&r'rGszCABlock._flash_attentioncCs*||dd|j}|jdd}||S)z=Attention matrix multiplication with depth-wise convolutions.r+) transposerBsoftmax)r#rMrNrOattnr&r&r'rHs zCABlock._normal_attentionr(r)r*c Cs|jdd}|||}|jddd\}}}|jdkr#d}d}nd}d }t|||jd }t|||jd }t|||jd }tjj j |d d}tjj j |d d}| |||} t| |fd |jit t |jdkrkd dgngd|} ttj|| S)aForward pass for MDTA attention. 1. Apply depth-wise convolutions to Q, K, V 2. Reshape Q, K, V for multi-head attention 3. Compute attention matrix using flash or normal attention 4. Reshape and project out attention outputrNrrr+z b (head c) h w -> b head c (h w)z b head c (h w) -> b (head c) h wz$b (head c) d h w -> b head c (d h w)z$b head c (d h w) -> b (head c) d h w)headrRrVhw)drWrX)shaperDrCr,r rr:r-r? functional normalizerFdictziprr.r") r#r(r rCrMrNrOZqkv_to_multiheadZmultihead_to_qkvrPr&r&r'r3s. "zCABlock.forward)F)r r r:r rrr;rr4) r5r6r7r8rrErGrHr3r9r&r&r$r'r Ns!) __future__rtypingrr-torch.nnr?torch.nn.functionalr[r/"monai.networks.blocks.convolutionsr monai.utilsrr___all__Moduler r r&r&r&r's    3