o %ݫi@sdZddlZddlmZddlmZmZddlmZee Z Gdddej j Z Gdd d ej j ZGd d d ej j ZGd d d ej j ZGdddej j ZGdddej j ZdS)zbLibrary implementing complex-valued recurrent neural networks. Authors * Titouan Parcollet 2020 N)CLinear) CBatchNorm CLayerNorm) get_loggercsJeZdZdZ       dfdd Zd d Zdd dZddZZS)CLSTMaThis function implements a complex-valued LSTM. Input format is (batch, time, fea) or (batch, time, fea, channel). In the latter shape, the two last dimensions will be merged: (batch, time, fea * channel) Arguments --------- hidden_size : int Number of output neurons (i.e, the dimensionality of the output). Specified value is in term of complex-valued neurons. Thus, the output is 2*hidden_size. input_shape : tuple The expected shape of the input. num_layers : int, optional Number of layers to employ in the RNN architecture (default 1). bias: bool, optional If True, the additive bias b is adopted (default True). dropout : float, optional It is the dropout factor (must be between 0 and 1) (default 0.0). bidirectional : bool, optional If True, a bidirectional model that scans the sequence both right-to-left and left-to-right is used (default False). return_hidden : bool, optional It True, the function returns the last hidden layer. init_criterion : str , optional (glorot, he). This parameter controls the initialization criterion of the weights. It is combined with weights_init to build the initialization method of the complex-valued weights (default "glorot"). weight_init : str, optional (complex, unitary). This parameter defines the initialization procedure of the complex-valued weights (default "complex"). "complex" will generate random complex-valued weights following the init_criterion and the complex polar form. "unitary" will normalize the weights to lie on the unit circle. More details in: "Deep Complex Networks", Trabelsi C. et al. Example ------- >>> inp_tensor = torch.rand([10, 16, 40]) >>> rnn = CLSTM(hidden_size=16, input_shape=inp_tensor.shape) >>> out_tensor = rnn(inp_tensor) >>> torch.Size([10, 16, 32]) TFglorotcomplexc st|d|_||_||_||_||_d|_||_||_ | |_ t |dkr+d|_t t |dd|_|d|_||_dSNFTr)super__init__ hidden_size num_layersbiasdropout bidirectionalreshape return_hiddeninit_criterion weight_initlentorchprodtensorfea_dim batch_size _init_layersrnn) selfr input_shaperrrrrrr __class__[/home/ubuntu/.local/lib/python3.10/site-packages/speechbrain/nnet/complex_networks/c_RNN.pyrCs   zCLSTM.__init__c Csntjg}|j}t|jD]&}t||j|j|j|j |j |j |j d}| ||j r1|jd}q|j}q|S)z Initializes the layers of the ComplexLSTM. Returns ------- rnn : ModuleList The list of CLSTM_Layers. )rrrrr )rnn ModuleListrranger CLSTM_Layerrrrrrrappendr!r current_dimirnn_layr%r%r&rcs$   zCLSTM._init_layersNcC^|jr|jdkr||jd|jd|jd|jd}|j||d\}}|jr-||fS|S)asReturns the output of the CLSTM. Arguments --------- x : torch.Tensor Input tensor. hx : torch.Tensor The hidden layer. Returns ------- output : torch.Tensor The output tensor. hh : torch.Tensor If return_hidden, the second tensor is hidden states. rrr r hxrndimshape _forward_rnnrr!xr3outputhhr%r%r&forward *z CLSTM.forwardcCg}|dur|jr||j|jd|j}t|jD]%\}}|dur+||||d}n||dd}||dddddfqtj |dd}|jr^||j dd|j d|j}||fS| dd}||fS)aXReturns the output of the CLSTM. Arguments --------- x : torch.Tensor Input tensor. hx : torch.Tensor The hidden layer. Returns ------- x : torch.Tensor The output tensor. h : torch.Tensor The hidden states for each step. Nr r2rdimr rrrrr enumerater r+rstackr6 transposer!r9r3hr.r/r%r%r&r7s"   zCLSTM._forward_rnn)rTrFFr r N __name__ __module__ __qualname____doc__rrr<r7 __classcell__r%r%r#r&rs3 !rcsTeZdZdZ    dfdd Zdd d Zd d Zd dZddZddZ Z S)r*aThis function implements complex-valued LSTM layer. Arguments --------- input_size : int Feature dimensionality of the input tensors (in term of real values). hidden_size : int Number of output values (in term of real values). num_layers : int, optional Number of layers to employ in the RNN architecture (default 1). batch_size : int Batch size of the input tensors. dropout : float, optional It is the dropout factor (must be between 0 and 1) (default 0.0). bidirectional : bool, optional If True, a bidirectional model that scans the sequence both right-to-left and left-to-right is used (default False). init_criterion : str, optional (glorot, he). This parameter controls the initialization criterion of the weights. It is combined with weights_init to build the initialization method of the complex-valued weights (default "glorot"). weight_init : str, optional (complex, unitary). This parameter defines the initialization procedure of the complex-valued weights (default "complex"). "complex" will generate random complex-valued weights following the init_criterion and the complex polar form. "unitary" will normalize the weights to lie on the unit circle. More details in: "Deep Complex Networks", Trabelsi C. et al. rFr r c stt|d|_t||_||_||_||_||_||_ t |j|jdd|j |jd|_ t |jd|jdd|j |jd|_ |jrK|jd|_| dtd|jd||jtjj|jdd|_td g|_dS) Nr r1Tr" n_neuronsrrrh_initrFpinplace?)rrintr input_sizerrrrrrwuregister_bufferrzeros _init_dropr'Dropoutdroprfloat drop_mask_te) r!rWrrrrrrrr#r%r&rs8    zCLSTM_Layer.__init__NcCs|jr|d}tj||gdd}||||}|dur&|||}n|||j}|jrG|jddd\}}|d}tj||gdd}|S)a/Returns the output of the CRNN_layer. Arguments --------- x : torch.Tensor Linearly transformed input. hx : torch.Tensor Hidden layer. Returns ------- h : torch.Tensor The hidden states for each step. rrr@Nr ) rfliprcat_change_batch_sizerX_complexlstm_cellrQchunkr!r9r3x_fliprXrGh_fh_br%r%r&r<s    zCLSTM_Layer.forwardc Csg}|j}||}t|jdD]a}|dd|f||}|dd\}} } } } } }}ttj|| gdd}ttj| | gdd}ttj| | gdd}|t tj||gdd|||}|t |}| |qtj |dd}|S)7Returns the hidden states for each time step. Arguments --------- w : torch.Tensor Linearly transformed input. ht : torch.Tensor Hidden layer. Returns ------- h : torch.Tensor The hidden states for each step. rNr?r@) rQ_sample_drop_maskr)r6rYrersigmoidrbtanhr+rD)r!rXhthiddensct drop_maskkgatesitritiftrftiotrotictrctiitftotrGr%r%r&rdEs"  zCLSTM_Layer._complexlstm_cellc CZtjj|jdd|_tdg|_d|_d|_ | d|t |j|j dj dS zwInitializes the recurrent dropout operation. To speed it up, the dropout masks are sampled in advance. FrRrUi>r drop_masksr Nrr'r]rr^rr_r` N_drop_masks drop_mask_cntrZonesrdatar!rr%r%r&r\pzCLSTM_Layer._init_dropcC|jr6|j|j|jkr!d|_|tj|j|jd|jdj |_ |j |j|j|j}|j|j|_|S|j |j|_ |j }|Sz,Selects one of the pre-defined dropout masksrr )device trainingrrrr^rrrrrrr`tor!rXrrr%r%r&rl$zCLSTM_Layer._sample_drop_maskcCL|j|jdkr"|jd|_|jr$|t|j|jdj|_ dSdSdSThis function changes the batch size when it is different from the one detected in the initialization method. This might happen in the case of multi-gpu or when we have different batch sizes in train and test. We also update the h_int and drop masks. rr N rr6rr^rrrrrrr!r9r%r%r&rc zCLSTM_Layer._change_batch_size)rFr r rH) rJrKrLrMrr<rdr\rlrcrNr%r%r#r&r*s% 3'+r*csLeZdZdZ        dfd d Zd d ZdddZddZZS)CRNNa@This function implements a vanilla complex-valued RNN. Input format is (batch, time, fea) or (batch, time, fea, channel). In the latter shape, the two last dimensions will be merged: (batch, time, fea * channel) Arguments --------- hidden_size : int Number of output neurons (i.e, the dimensionality of the output). Specified value is in term of complex-valued neurons. Thus, the output is 2*hidden_size. input_shape : tuple The expected shape of the input. nonlinearity : str, optional Type of nonlinearity (tanh, relu) (default "tanh"). num_layers : int, optional Number of layers to employ in the RNN architecture (default 1). bias : bool, optional If True, the additive bias b is adopted (default True). dropout : float, optional It is the dropout factor (must be between 0 and 1) (default 0.0). bidirectional : bool, optional If True, a bidirectional model that scans the sequence both right-to-left and left-to-right is used (default False). return_hidden : bool, optional It True, the function returns the last hidden layer (default False). init_criterion : str , optional (glorot, he). This parameter controls the initialization criterion of the weights. It is combined with weights_init to build the initialization method of the complex-valued weights (default "glorot"). weight_init : str, optional (complex, unitary). This parameter defines the initialization procedure of the complex-valued weights (default "complex"). "complex" will generate random complex-valued weights following the init_criterion and the complex polar form. "unitary" will normalize the weights to lie on the unit circle. More details in: "Deep Complex Networks", Trabelsi C. et al. Example ------- >>> inp_tensor = torch.rand([10, 16, 30]) >>> rnn = CRNN(hidden_size=16, input_shape=inp_tensor.shape) >>> out_tensor = rnn(inp_tensor) >>> torch.Size([10, 16, 32]) rnrTrFr r c st|d|_||_||_||_||_||_d|_||_ | |_ | |_ t |dkr.d|_t t |dd|_|d|_||_dSr )rrr nonlinearityrrrrrrrrrrrrrrrr ) r!rr"rrrrrrrrr#r%r&rs   z CRNN.__init__c Csrtjg}|j}t|jD](}t||j|j|j|j |j |j |j |j d }|||j r3|jd}q|j}q|S)z Initializes the layers of the CRNN. Returns ------- rnn : ModuleList The list of CRNN_Layers. )rrrrrr )rr'r(rr)r CRNN_Layerrrrrrrrr+r,r%r%r&rs&   zCRNN._init_layersNcCr0)aReturns the output of the vanilla CRNN. Arguments --------- x : torch.Tensor Input tensor. hx : torch.Tensor Hidden layers. Returns ------- output : torch.Tensor The outputs of the CliGRU. hh : torch.Tensor If return_hidden, also returns the hidden states for each step. r1rrr r r2r4r8r%r%r&r<r=z CRNN.forwardcCr>)a_Returns the output of the vanilla CRNN. Arguments --------- x : torch.Tensor Input tensor. hx : torch.Tensor The hidden layer. Returns ------- x : torch.Tensor The output tensor. h : torch.Tensor The hidden states for each step. Nr r2r?rr@rrBrFr%r%r&r7=s"   zCRNN._forward_rnn)rnrTrFFr r rHrIr%r%r#r&rs5" "rcsVeZdZdZ     dfdd Zdd d Zd d ZddZddZddZ Z S)raThis function implements complex-valued recurrent layer. Arguments --------- input_size : int Feature dimensionality of the input tensors (in term of real values). hidden_size : int Number of output values (in term of real values). num_layers : int, optional Number of layers to employ in the RNN architecture (default 1). batch_size : int Batch size of the input tensors. dropout : float, optional It is the dropout factor (must be between 0 and 1) (default 0.0). nonlinearity : str, optional Type of nonlinearity (tanh, relu) (default "tanh"). bidirectional : bool, optional If True, a bidirectional model that scans the sequence both right-to-left and left-to-right is used (default False). init_criterion : str , optional (glorot, he). This parameter controls the initialization criterion of the weights. It is combined with weights_init to build the initialization method of the complex-valued weights (default "glorot"). weight_init : str, optional (complex, unitary). This parameter defines the initialization procedure of the complex-valued weights (default "complex"). "complex" will generate random complex-valued weights following the init_criterion and the complex polar form. "unitary" will normalize the weights to lie on the unit circle. More details in: "Deep Complex Networks", Trabelsi C. et al. rrnFr r c stt|d|_t||_||_||_||_||_| |_ t |j|jd|j |jd|_ t |jd|jd|j |jd|_ |jrG|jd|_| dtd|jd||jtjj|jdd|_tdg|_|dkrytj|_dStj|_dS) Nr FrOrQrrRrUrn)rrrVrrWrrrrrrrXrYrZrr[r\r'r]r^rr_r`TanhactReLU) r!rWrrrrrrrrr#r%r&rs>    zCRNN_Layer.__init__NcCs|jr|d}tj||gdd}||}|dur!|||}n|||j}|jrB|jddd\}}|d}tj||gdd}|S)a%Returns the output of the CRNN_layer. Arguments --------- x : torch.Tensor Input tensor. hx : torch.Tensor The hidden layer. Returns ------- h : torch.Tensor The hidden states for each step. rrr@Nr )rrarrbrX_complexrnn_cellrQrerfr%r%r&r<s   zCRNN_Layer.forwardcCsfg}||}t|jdD]}|dd|f||}|||}||qtj|dd}|S)a;Returns the hidden states for each time step. Arguments --------- w : torch.Tensor Linearly transformed input. ht : torch.Tensor The hidden layer. Returns ------- h : torch.Tensor The hidden states for each step. rNr@)rlr)r6rYrr+rrD)r!rXrorprrrsatrGr%r%r&rs  zCRNN_Layer._complexrnn_cellc Crrrrr%r%r&r\ rzCRNN_Layer._init_dropcCrrrrr%r%r&rlrzCRNN_Layer._sample_drop_maskcCrrrrr%r%r&rc2rzCRNN_Layer._change_batch_size)rrnFr r rH) rJrKrLrMrr<rr\rlrcrNr%r%r#r&rgs' 9(rcsNeZdZdZ         dfd d Zd d ZdddZddZZS)CLiGRUaN This function implements a complex-valued Light GRU (liGRU). Ligru is single-gate GRU model based on batch-norm + relu activations + recurrent dropout. For more info see: "M. Ravanelli, P. Brakel, M. Omologo, Y. Bengio, Light Gated Recurrent Units for Speech Recognition, in IEEE Transactions on Emerging Topics in Computational Intelligence, 2018" (https://arxiv.org/abs/1803.10225) To speed it up, it is compiled with the torch just-in-time compiler (jit) right before using it. It accepts in input tensors formatted as (batch, time, fea). In the case of 4d inputs like (batch, time, fea, channel) the tensor is flattened as (batch, time, fea*channel). Arguments --------- hidden_size : int Number of output neurons (i.e, the dimensionality of the output). Specified value is in term of complex-valued neurons. Thus, the output is 2*hidden_size. input_shape : tuple The expected size of the input. nonlinearity : str Type of nonlinearity (tanh, relu). normalization : str Type of normalization for the ligru model (batchnorm, layernorm). Every string different from batchnorm and layernorm will result in no normalization. num_layers : int Number of layers to employ in the RNN architecture. bias : bool If True, the additive bias b is adopted. dropout : float It is the dropout factor (must be between 0 and 1). bidirectional : bool If True, a bidirectional model that scans the sequence both right-to-left and left-to-right is used. return_hidden : bool If True, the function returns the last hidden layer. init_criterion : str , optional (glorot, he). This parameter controls the initialization criterion of the weights. It is combined with weights_init to build the initialization method of the complex-valued weights (default "glorot"). weight_init : str, optional (complex, unitary). This parameter defines the initialization procedure of the complex-valued weights (default "complex"). "complex" will generate random complex-valued weights following the init_criterion and the complex polar form. "unitary" will normalize the weights to lie on the unit circle. More details in: "Deep Complex Networks", Trabelsi C. et al. Example ------- >>> inp_tensor = torch.rand([10, 16, 30]) >>> rnn = CLiGRU(input_shape=inp_tensor.shape, hidden_size=16) >>> out_tensor = rnn(inp_tensor) >>> torch.Size([4, 10, 5]) relu batchnormrTrFr r c st|d|_||_||_||_||_||_||_d|_ | |_ | |_ | |_ t |dkr1d|_ tt|dd|_|d|_||_dSr )rrrrr normalizationrrrrrrrrrrrrrrr ) r!rr"rrrrrrrrrr#r%r&rs"    zCLiGRU.__init__c Csvtjg}|j}t|jD]*}t||j|j|j|j |j |j |j |j |jd }|||j r5|jd}q|j}q|S)zInitializes the layers of the liGRU. Returns ------- rnn : ModuleList The list of CLiGRU_Layers. )rrrrrrr )rr'r(rr)r CLiGRU_Layerrrrrrrrrr+r,r%r%r&rs(   zCLiGRU._init_layersNcCr0)aReturns the output of the CliGRU. Arguments --------- x : torch.Tensor Input tensor. hx : torch.Tensor Hidden layers. Returns ------- output : torch.Tensor The outputs of the CliGRU. hh : torch.Tensor If return_hidden, also returns the hidden states for each step. r1rrr r r2)rr5r6_forward_ligrurr8r%r%r&r<s *zCLiGRU.forwardcCr>)aYReturns the output of the CliGRU. Arguments --------- x : torch.Tensor Input tensor. hx : torch.Tensor The hidden layer. Returns ------- x : torch.Tensor The output tensor. h : torch.Tensor The hidden states for each step. Nr r2r?rr@rrB)r!r9r3rGr. ligru_layr%r%r&rs"   zCLiGRU._forward_ligru) rrrTrFFr r rH) rJrKrLrMrrr<rrNr%r%r#r&rBsD" !rcsXeZdZdZ      dfdd Zdd d Zd dZddZddZddZ Z S)ra This function implements complex-valued Light-Gated Recurrent Unit layer. Arguments --------- input_size : int Feature dimensionality of the input tensors. hidden_size : int Number of output values. num_layers : int Number of layers to employ in the RNN architecture. batch_size : int Batch size of the input tensors. dropout : float It is the dropout factor (must be between 0 and 1). nonlinearity : str Type of nonlinearity (tanh, relu). normalization : str Type of normalization (batchnorm, layernorm). Every string different from batchnorm and layernorm will result in no normalization. bidirectional : bool If True, a bidirectional model that scans the sequence both right-to-left and left-to-right is used. init_criterion : str , optional (glorot, he). This parameter controls the initialization criterion of the weights. It is combined with weights_init to build the initialization method of the complex-valued weights (default "glorot"). weight_init : str, optional (complex, unitary). This parameter defines the initialization procedure of the complex-valued weights (default "complex"). "complex" will generate random complex-valued weights following the init_criterion and the complex polar form. "unitary" will normalize the weights to lie on the unit circle. More details in: "Deep Complex Networks", Trabelsi C. et al. rrrFr r c stt|d|_t||_||_||_||_| |_| |_ ||_ ||_ t |j|jdd|j |jd|_ t |jd|jdd|j |jd|_|jrQ|jd|_d|_|j dkrgt|dddd|_d|_n|j d kryt|ddd |_d|_n t|ddd |_d|_|d td |jd||jtjj|jdd |_tdg|_|j dkrtj|_dStj|_dS)Nr FrOrr?g?)rWrAmomentumT layernorm)rWrArQrrRrUrn) rrrVrrWrrrrrrrrrXrY normalizernormrrZrr[r\r'r]r^rr_r`rrr) r!rWrrrrrrrrrr#r%r&r5sX        zCLiGRU_Layer.__init__Nc Cs|jr|d}tj||gdd}||||}|jrB|||j d|j d|j d}||j d|j d|j d}|durM| ||}n| ||j }|jrn|j ddd\}}|d}tj||gdd}|S)a*Returns the output of the Complex liGRU layer. Arguments --------- x : torch.Tensor Input tensor. hx : torch.Tensor Hidden layer. Returns ------- h : torch.Tensor The hidden states for each step. rrr@r N) rrarrbrcrXrrrr6_complex_ligru_cellrQre) r!r9r3rgrXw_bnrGrhrir%r%r&r<s    (  zCLiGRU_Layer.forwardcCsg}||}t|jdD]F}|dd|f||}|dd\}}} } tj||gdd} tj| | gdd} t| } || |} | |d| | }| |qtj |dd}|S)rjrNr1r?r@) rlr)r6rYrerrbrmrr+rD)r!rXrorprrrsrtatratiztrztirzthcandrGr%r%r&rs   z CLiGRU_Layer._complex_ligru_cellc Crrrrr%r%r&r\rzCLiGRU_Layer._init_dropcCrrrrr%r%r&rlrzCLiGRU_Layer._sample_drop_maskcCsH|j|jdkr |jd|_|jr"|t|j|jj|_ dSdSdS)rrNrrr%r%r&rcs zCLiGRU_Layer._change_batch_size)rrrFr r rH) rJrKrLrMrr<rr\rlrcrNr%r%r#r&rs, N-$r)rMr*speechbrain.nnet.complex_networks.c_linearr1speechbrain.nnet.complex_networks.c_normalizationrrspeechbrain.utils.loggerrrJloggerr'Modulerr*rrrrr%r%r%r&s"  9_?\M