o & iH'@sbddlmZddlmZmZddlZddlmZddgZGdddejZ GdddejjZ dS) ) annotations)SequenceTupleN)nnVectorQuantizer EMAQuantizercsVeZdZdZ     d%d&fdd Zd'ddZd(ddZd)d!d"Zd'd#d$ZZ S)*ra9 Vector Quantization module using Exponential Moving Average (EMA) to learn the codebook parameters based on Neural Discrete Representation Learning by Oord et al. (https://arxiv.org/abs/1711.00937) and the official implementation that can be found at https://github.com/deepmind/sonnet/blob/v2/sonnet/src/nets/vqvae.py#L148 and commit 58d9a2746493717a7c9252938da7efa6006f3739. This module is not compatible with TorchScript while working in a Distributed Data Parallelism Module. This is due to lack of TorchScript support for torch.distributed module as per https://github.com/pytorch/pytorch/issues/41353 on 22/10/2022. If you want to TorchScript your model, please turn set `ddp_sync` to False. Args: spatial_dims: number of spatial dimensions of the input. num_embeddings: number of atomic elements in the codebook. embedding_dim: number of channels of the input and atomic elements. commitment_cost: scaling factor of the MSE loss between input and its quantized version. Defaults to 0.25. decay: EMA decay. Defaults to 0.99. epsilon: epsilon value. Defaults to 1e-5. embedding_init: initialization method for the codebook. Defaults to "normal". ddp_sync: whether to synchronize the codebook across processes. Defaults to True. ?Gz?h㈵>normalT spatial_dimsintnum_embeddings embedding_dimcommitment_costfloatdecayepsilonembedding_initstrddp_syncboolc st||_||_||_|jdvsJtd|dtj|j|j|_ |dkr,n|dkr=tjj j |j j j dddd |j j _||_|d t|j|d |j j j ||||_||_||_d gttd |jd dg|_d |jdgttd|jd|_dS)N)zMEMAQuantizer only supports 4D and 5D tensor inputs but received spatial dims .r kaiming_uniformfan_inlinear)mode nonlinearityFema_cluster_sizeema_wrr)super__init__r rr ValueErrortorchr Embedding embeddinginitkaiming_uniform_weightdata requires_gradrregister_bufferzerosclonerrrlistrangeflatten_permutationquantization_permutation) selfr rrrrrrr __class__h/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/layers/vector_quantizer.pyr$,s2   " zEMAQuantizer.__init__inputs torch.Tensorreturn/Tuple[torch.Tensor, torch.Tensor, torch.Tensor]cCstjdddbt|j}|d=|}||jd|j }|dj ddd|j j dj d dddt||j j }tj| dd d}tjj||j}||}|||fWd S1slwYd S) a Given an input it projects it to the quantized space and returns additional tensors needed for EMA loss. Args: inputs: Encoding space tensors of shape [B, C, H, W, D]. Returns: torch.Tensor: Flatten version of the input of shape [B*H*W*D, C]. torch.Tensor: One-hot representation of the quantization indices of shape [B*H*W*D, self.num_embeddings]. torch.Tensor: Quantization indices of shape [B,H,W,D,1] cudaFenabledr"rT)dimkeepdimr)rBN)r&autocastr1shaperpermuter3 contiguousviewrsumr(r+tmmmaxr functionalone_hotr)r5r:Zencoding_indices_view flat_input distancesencoding_indices encodingsr8r8r9quantizeZs   $zEMAQuantizer.quantizeembedding_indicescCsLtjddd|||j}|WdS1swYdS)a Given encoding indices of shape [B,D,H,W,1] embeds them in the quantized space [B, D, H, W, self.embedding_dim] and reshapes them to [B, self.embedding_dim, D, H, W] to be fed to the decoder. Args: embedding_indices: Tensor in channel last format which holds indices referencing atomic elements from self.embedding Returns: torch.Tensor: Quantize space representation of encoding_indices in channel first format. r>Fr?N)r&rDr(rFr4rG)r5rTr(r8r8r9embeds  $zEMAQuantizer.embed encodings_sumdwNonecCsF|jr tjr tjj|tjjjdtjj|tjjjddS dS)a' TorchScript does not support torch.distributed.all_reduce. This function is a bypassing trick based on the example: https://pytorch.org/docs/stable/generated/torch.jit.unused.html#torch.jit.unused Args: encodings_sum: The summation of one hot representation of what encoding was used for each position. dw: The multiplication of the one hot representation of what encoding was used for each position with the flattened input. Returns: None )tensoropN)rr& distributedis_initialized all_reduceReduceOpSUM)r5rVrWr8r8r9distributed_synchronizationsz(EMAQuantizer.distributed_synchronizationc Cs6||\}}}||}|jrte|d}t||}|jr+| |||j j |j t|d|j |j }|j |j||j|j|} |jj |j t|d|j |jjj |j| dWdn1s|wY|jtjj||} |||}|| |fS)Nrr")rSrUtrainingr&no_gradrIrKrJrr`r r,mul_radd_mulrrr!r(r+copy_ unsqueezerrrMmse_lossdetach) r5r:rOrRrQ quantizedrVrWnweightslossr8r8r9forwards"    $  $ zEMAQuantizer.forward)rr r r T)r r rr rr rrrrrrrrrr)r:r;r<r=rTr;r<r;)rVr;rWr;r<rX) __name__ __module__ __qualname____doc__r$rSrUr`rn __classcell__r8r8r6r9rs . & cs@eZdZdZdfdd Zdd d Zdd dZdddZZS)raU Vector Quantization wrapper that is needed as a workaround for the AMP to isolate the non fp16 compatible parts of the quantization in their own class. Args: quantizer (torch.nn.Module): Quantizer module that needs to return its quantized representation, loss and index based quantized representation. quantizerrcs t||_td|_dS)Nr")r#r$rur&rand perplexity)r5rur6r8r9r$s zVectorQuantizer.__init__r:r;r<!Tuple[torch.Tensor, torch.Tensor]c Csd||\}}}tj||jj|jjd|}tt|t |d |_ ||fS)N)binsrLg|=) rur&histcrrdivnumelexprIlogrw)r5r:rjrmrQZ avg_probsr8r8r9rns "zVectorQuantizer.forwardrTcCs|jj|dS)N)rT)rurU)r5rTr8r8r9rUszVectorQuantizer.embedrRcCs||}|d}|S)Nr)ru)r5rRoutputrQr8r8r9rSs zVectorQuantizer.quantize)rur)r:r;r<rxro)rRr;r<r;) rprqrrrsr$rnrUrSrtr8r8r6r9rs   ) __future__rtypingrrr&r__all__Modulerrr8r8r8r9s  0