o - iG@sddlmZddlmZddlmZddlZddlmZddl m Z ddl m Z ddl mZmZddlmZd gZGd d d ejZGd d d ejZGdddejZGdd d ejZdS)) annotations)Sequence)TupleN) Convolution)Act) EMAQuantizerVectorQuantizer)ensure_tuple_repVQVAEcs4eZdZdZejddfdfdd ZddZZS)VQVAEResidualUnita Implementation of the ResidualLayer used in the VQVAE network as originally used in Morphology-preserving Autoregressive 3D Generative Modelling of the Brain by Tudosiu et al. (https://arxiv.org/pdf/2209.03177.pdf). The original implementation that can be found at https://github.com/AmigoLab/SynthAnatomy/blob/main/src/networks/vqvae/baseline.py#L150. Args: spatial_dims: number of spatial spatial_dims of the input data. in_channels: number of input channels. num_res_channels: number of channels in the residual layers. act: activation type and arguments. Defaults to RELU. dropout: dropout ratio. Defaults to no dropout. bias: whether to have a bias term. Defaults to True. T spatial_dimsint in_channelsnum_res_channelsacttuple | str | NonedropoutfloatbiasboolreturnNonec srt||_||_||_||_||_||_t|j|j|jd|j|j|jd|_ t|j|j|j|jdd|_ dS)NDA)r r out_channels adn_orderingrrrT)r rrr conv_only) super__init__r rrrrrrconv1conv2)selfr rrrrr __class__[/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/vqvae.pyr-s.   zVQVAEResidualUnit.__init__cCs tjj||||dS)NT)torchnn functionalrelur r)r!xr$r$r%forwardQs zVQVAEResidualUnit.forward)r rrrrrrrrrrrrr) __name__ __module__ __qualname____doc__rRELUrr+ __classcell__r$r$r"r%r s$r cs,eZdZdZdfdd ZdddZZS)Encodera Encoder module for VQ-VAE. Args: spatial_dims: number of spatial spatial_dims. in_channels: number of input channels. out_channels: number of channels in the latent space (embedding_dim). channels: sequence containing the number of channels at each level of the encoder. num_res_layers: number of sequential residual layers at each level. num_res_channels: number of channels in the residual layers at each level. downsample_parameters: A Tuple of Tuples for defining the downsampling convolutions. Each Tuple should hold the following information stride (int), kernel_size (int), dilation (int) and padding (int). dropout: dropout ratio. act: activation type and arguments. r rrrchannels Sequence[int]num_res_layersrdownsample_parameters#Sequence[Tuple[int, int, int, int]]rrrrrrc sNt||_||_||_||_||_||_||_||_ | |_ g} t t |jD]]} | t|j| dkr7|jn|j| d|j| |j| d|j| dd|j | dkrVdn|j d|j| d|j| dd t |jD]} | t|j|j| |j| |j |j dqnq)| t|j|jt |jd|jddddd t| |_dS) Nrr) r rrstrides kernel_sizerrr dropout_dimdilationpaddingr rrrrTr rrr;r<r?r)rrr rrr3r5rr6rrrangelenappendrr r' ModuleListblocks) r!r rrr3r5rr6rrrFi_r"r$r%rfsd       zEncoder.__init__r* torch.TensorcC|jD]}||}q|SNrFr!r*blockr$r$r%r+  zEncoder.forward)r rrrrrr3r4r5rrr4r6r7rrrrrrr*rIrrIr,r-r.r/rr+r1r$r$r"r%r2UsCr2cs,eZdZdZdfdd ZdddZZS)DecoderaW Decoder module for VQ-VAE. Args: spatial_dims: number of spatial spatial_dims. in_channels: number of channels in the latent space (embedding_dim). out_channels: number of output channels. channels: sequence containing the number of channels at each level of the decoder. num_res_layers: number of sequential residual layers at each level. num_res_channels: number of channels in the residual layers at each level. upsample_parameters: A Tuple of Tuples for defining the upsampling convolutions. Each Tuple should hold the following information stride (int), kernel_size (int), dilation (int), padding (int), output_padding (int). dropout: dropout ratio. act: activation type and arguments. output_act: activation type and arguments for the output. r rrrr3r4r5rupsample_parameters(Sequence[Tuple[int, int, int, int, int]]rrrr output_actrrc st||_||_||_||_||_||_||_||_ | |_ | |_ t t |j} g} | t|j|j| ddddddt t |j} tt|jD]r}t|jD]}| t|j| || ||j |j dqS| t|j| ||t|jdkr}|jn| |d|j|d|j|dd|j |t|jdkr|j ndd|j|d|t|jdkd|j|d|j|d d qL|j r| t|j t| |_dS) Nrr8r:TrAr@rr9)r rrr;r<rrrnormr>r is_transposedr?output_padding)rrr rrr3r5rrSrrrUlistreversedrDrrBrCr rr'rErF)r!r rrr3r5rrSrrrUZreversed_num_channelsrFZreversed_num_res_channelsrGrHr"r$r%rsr   "     zDecoder.__init__r*rIcCrJrKrLrMr$r$r%r+rOzDecoder.forward)r rrrrrr3r4r5rrr4rSrTrrrrrUrrrrPrQr$r$r"r%rRsMrRcseZdZdZddddddddd d d d ejd ddfdGfd,d- ZdHd1d2ZdId5d6ZdJd8d9Z dHd:d;Z dKd=d>Z dLd?d@Z dMdBdCZ dNdEdFZZS)Or a Vector-Quantised Variational Autoencoder (VQ-VAE) used in Morphology-preserving Autoregressive 3D Generative Modelling of the Brain by Tudosiu et al. (https://arxiv.org/pdf/2209.03177.pdf) The original implementation can be found at https://github.com/AmigoLab/SynthAnatomy/blob/main/src/networks/vqvae/baseline.py#L163/ Args: spatial_dims: number of spatial spatial_dims. in_channels: number of input channels. out_channels: number of output channels. downsample_parameters: A Tuple of Tuples for defining the downsampling convolutions. Each Tuple should hold the following information stride (int), kernel_size (int), dilation (int) and padding (int). upsample_parameters: A Tuple of Tuples for defining the upsampling convolutions. Each Tuple should hold the following information stride (int), kernel_size (int), dilation (int), padding (int), output_padding (int). num_res_layers: number of sequential residual layers at each level. channels: number of channels at each level. num_res_channels: number of channels in the residual layers at each level. num_embeddings: VectorQuantization number of atomic elements in the codebook. embedding_dim: VectorQuantization number of channels of the input and atomic elements. commitment_cost: VectorQuantization commitment_cost. decay: VectorQuantization decay. epsilon: VectorQuantization epsilon. act: activation type and arguments. dropout: dropout ratio. output_act: activation type and arguments for the output. ddp_sync: whether to synchronize the codebook across processes. use_checkpointing if True, use activation checkpointing to save memory. )`r\r:)r9rVr8r8r^r^)r9rVr8r8rr_r_ @normalg?g?gh㈵>r NTFr rrrr3r4r5rSequence[int] | intr6?Sequence[Tuple[int, int, int, int]] | Tuple[int, int, int, int]rSISequence[Tuple[int, int, int, int, int]] | Tuple[int, int, int, int, int]num_embeddings embedding_dimembedding_initstrcommitment_costrdecayepsilonrrrrUddp_syncruse_checkpointingc st||_||_||_||_| |_| |_||_t |t r&t |t |}t |t |kr2t dtdd|DrC|ft |}n|}tdd|DrV|ft |}n|}tdd|Dset dtdd|Dsrt d|D] }t |d krt d qt|D] }t |d krt d qt |t |krt d t |t |krt d||_||_t||| ||||||d |_t|| ||||||||d |_tt|| | | | || |dd|_dS)Nzk`num_res_channels` should be a single integer or a tuple of integers with the same length as `num_channls`.cs|]}t|tVqdSrK isinstancer.0valuesr$r$r% cz!VQVAE.__init__..csrorKrprrr$r$r%ruhrvcs"|] }tdd|DVqdS)csrorKrprsvaluer$r$r%rumrv+VQVAE.__init__...NallrsZsub_itemr$r$r%rum zQ`downsample_parameters` should be a single tuple of integer or a tuple of tuples.csrw)csrorKrprxr$r$r%ruqrvrzNr{r}r$r$r%ruqr~zO`upsample_parameters` should be a single tuple of integer or a tuple of tuples.rVzD`downsample_parameters` should be a tuple of tuples with 4 integers.zB`upsample_parameters` should be a tuple of tuples with 5 integers.z[`downsample_parameters` should be a tuple of tuples with the same length as `num_channels`.zY`upsample_parameters` should be a tuple of tuples with the same length as `num_channels`.) r rrr3r5rr6rr) r rrr3r5rrSrrrU)r rfrgrjrkrlrhrm quantizer)rrrrr r3rfrgrnrqrr rC ValueErrorr|r5rr2encoderrRdecoderrrr)r!r rrr3r5rr6rSrfrgrhrjrkrlrrrUrmrnZupsample_parameters_tupleZdownsample_parameters_tuple parameterr"r$r%r3s       zVQVAE.__init__imagesrIrcC.|jrtjjj|j|dd}|S||}|SNF) use_reentrant)rnr&utils checkpointr)r!routputr$r$r%encodes  z VQVAE.encode encodings!tuple[torch.Tensor, torch.Tensor]cCs||\}}||fSrKr)r!rZx_lossr*r$r$r%quantizeszVQVAE.quantize quantizationscCrr)rnr&rrr)r!rrr$r$r%decodes  z VQVAE.decodecCs|j|j|dS)N)r)rrr)r!rr$r$r%index_quantizeszVQVAE.index_quantizeembedding_indicescCs||j|SrK)rrembed)r!rr$r$r%decode_samplesszVQVAE.decode_samplescCs&|||\}}||}||fSrK)rrr)r!rrZquantization_lossesreconstructionr$r$r%r+s z VQVAE.forwardr*cCs||}||\}}|SrK)rr)r!r*zerHr$r$r%encode_stage_2_inputss zVQVAE.encode_stage_2_inputsrcCs||\}}||}|SrK)rr)r!rrrHimager$r$r%decode_stage_2_outputss zVQVAE.decode_stage_2_outputs)&r rrrrrr3r4r5rrrcr6rdrSrerfrrgrrhrirjrrkrrlrrrrrrUrrmrrnr)rrIrrI)rrIrr)rrIrrI)rrIrrI)rrIrrrP)rrIrrI)r,r-r.r/rr0rrrrrrr+rrr1r$r$r"r%r s6# |     ) __future__rcollections.abcrtypingrr&torch.nnr'monai.networks.blocksrmonai.networks.layersrZ&monai.networks.layers.vector_quantizerrr monai.utilsr __all__Moduler r2rRr r$r$r$r%s      9Ze