<reponame>sharod-mancity/speechbrain
##########################################################
# pytorch-kaldi v.0.1                                      
# Author: <NAME>, <NAME>
# Mila, University of Montreal
# October 2018
##########################################################

import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np
from distutils.util import strtobool

class LayerNorm(nn.Module):

    def __init__(self, features, eps=1e-6):
        super(LayerNorm,self).__init__()
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

def act_fun(act_type):

    if act_type=="relu":
        return nn.ReLU()
            
    if act_type=="tanh":
        return nn.Tanh()
            
    if act_type=="sigmoid":
        return nn.Sigmoid()
           
    if act_type=="leaky_relu":
        return nn.LeakyReLU(0.2)
            
    if act_type=="elu":
        return nn.ELU()
                     
    if act_type=="softmax":
        return nn.LogSoftmax(dim=1)
        
    if act_type=="linear":
        return nn.LeakyReLU(1)


class MLP(nn.Module):
    def __init__(self, options, inp_dim):
        super(MLP, self).__init__()
        
        self.input_dim=inp_dim
        self.dnn_lay=list(map(int, options['dnn_lay'].split(',')))
        self.dnn_drop=list(map(float, options['dnn_drop'].split(','))) 
        self.dnn_use_batchnorm=list(map(strtobool, options['dnn_use_batchnorm'].split(',')))
        self.dnn_use_laynorm=list(map(strtobool, options['dnn_use_laynorm'].split(','))) 
        self.dnn_use_laynorm_inp=strtobool(options['dnn_use_laynorm_inp'])
        self.dnn_use_batchnorm_inp=strtobool(options['dnn_use_batchnorm_inp'])
        self.dnn_act=options['dnn_act'].split(',')
        
       
        self.wx  = nn.ModuleList([])
        self.bn  = nn.ModuleList([])
        self.ln  = nn.ModuleList([])
        self.act = nn.ModuleList([])
        self.drop = nn.ModuleList([])
       
  
        # input layer normalization
        if self.dnn_use_laynorm_inp:
           self.ln0=LayerNorm(self.input_dim)
          
        # input batch normalization    
        if self.dnn_use_batchnorm_inp:
            self.bn0=nn.BatchNorm1d(self.input_dim,momentum=0.05)
           
        self.N_dnn_lay=len(self.dnn_lay)
             
        current_input=self.input_dim
        
        # Initialization of hidden layers
        
        for i in range(self.N_dnn_lay):
            
             # dropout
             self.drop.append(nn.Dropout(p=self.dnn_drop[i]))
             
             # activation
             self.act.append(act_fun(self.dnn_act[i]))
             
             
             add_bias=True
             
             # layer norm initialization
             self.ln.append(LayerNorm(self.dnn_lay[i]))
             self.bn.append(nn.BatchNorm1d(self.dnn_lay[i],momentum=0.05))
             
             if self.dnn_use_laynorm[i] or self.dnn_use_batchnorm[i]:
                 add_bias=False
             
             # Linear operations
             self.wx.append(nn.Linear(current_input, self.dnn_lay[i],bias=add_bias))
             
             # weight initialization
             self.wx[i].weight = torch.nn.Parameter(torch.Tensor(self.dnn_lay[i],current_input).uniform_(-np.sqrt(0.01/(current_input+self.dnn_lay[i])),np.sqrt(0.01/(current_input+self.dnn_lay[i]))))
             self.wx[i].bias = torch.nn.Parameter(torch.zeros(self.dnn_lay[i]))
             
             current_input=self.dnn_lay[i]
             
        self.out_dim=current_input
         
    def forward(self, x):
        
      # Applying Layer/Batch Norm
      if bool(self.dnn_use_laynorm_inp):
        x=self.ln0((x))
        
      if bool(self.dnn_use_batchnorm_inp):
        x=self.bn0((x))
        
      for i in range(self.N_dnn_lay):
           
          if self.dnn_use_laynorm[i] and not(self.dnn_use_batchnorm[i]):
           x = self.drop[i](self.act[i](self.ln[i](self.wx[i](x))))
          
          if self.dnn_use_batchnorm[i] and not(self.dnn_use_laynorm[i]):
           x = self.drop[i](self.act[i](self.bn[i](self.wx[i](x))))
           
          if self.dnn_use_batchnorm[i]==True and self.dnn_use_laynorm[i]==True:
           x = self.drop[i](self.act[i](self.bn[i](self.ln[i](self.wx[i](x)))))
          
          if self.dnn_use_batchnorm[i]==False and self.dnn_use_laynorm[i]==False:
           x = self.drop[i](self.act[i](self.wx[i](x)))
            
          
      return x

##########################################################
# pytorch-kaldi v.0.1                                      
# Author: <NAME>, <NAME>
# Mila, University of Montreal
# October 2018
##########################################################

import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np
from distutils.util import strtobool

class LayerNorm(nn.Module):

    def __init__(self, features, eps=1e-6):
        super(LayerNorm,self).__init__()
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

def act_fun(act_type):

    if act_type=="relu":
        return nn.ReLU()
            
    if act_type=="tanh":
        return nn.Tanh()
            
    if act_type=="sigmoid":
        return nn.Sigmoid()
           
    if act_type=="leaky_relu":
        return nn.LeakyReLU(0.2)
            
    if act_type=="elu":
        return nn.ELU()
                     
    if act_type=="softmax":
        return nn.LogSoftmax(dim=1)
        
    if act_type=="linear":
        return nn.LeakyReLU(1)


class MLP(nn.Module):
    def __init__(self, options, inp_dim):
        super(MLP, self).__init__()
        
        self.input_dim=inp_dim
        self.dnn_lay=list(map(int, options['dnn_lay'].split(',')))
        self.dnn_drop=list(map(float, options['dnn_drop'].split(','))) 
        self.dnn_use_batchnorm=list(map(strtobool, options['dnn_use_batchnorm'].split(',')))
        self.dnn_use_laynorm=list(map(strtobool, options['dnn_use_laynorm'].split(','))) 
        self.dnn_use_laynorm_inp=strtobool(options['dnn_use_laynorm_inp'])
        self.dnn_use_batchnorm_inp=strtobool(options['dnn_use_batchnorm_inp'])
        self.dnn_act=options['dnn_act'].split(',')
        
       
        self.wx  = nn.ModuleList([])
        self.bn  = nn.ModuleList([])
        self.ln  = nn.ModuleList([])
        self.act = nn.ModuleList([])
        self.drop = nn.ModuleList([])
       
  
        # input layer normalization
        if self.dnn_use_laynorm_inp:
           self.ln0=LayerNorm(self.input_dim)
          
        # input batch normalization    
        if self.dnn_use_batchnorm_inp:
            self.bn0=nn.BatchNorm1d(self.input_dim,momentum=0.05)
           
        self.N_dnn_lay=len(self.dnn_lay)
             
        current_input=self.input_dim
        
        # Initialization of hidden layers
        
        for i in range(self.N_dnn_lay):
            
             # dropout
             self.drop.append(nn.Dropout(p=self.dnn_drop[i]))
             
             # activation
             self.act.append(act_fun(self.dnn_act[i]))
             
             
             add_bias=True
             
             # layer norm initialization
             self.ln.append(LayerNorm(self.dnn_lay[i]))
             self.bn.append(nn.BatchNorm1d(self.dnn_lay[i],momentum=0.05))
             
             if self.dnn_use_laynorm[i] or self.dnn_use_batchnorm[i]:
                 add_bias=False
             
             # Linear operations
             self.wx.append(nn.Linear(current_input, self.dnn_lay[i],bias=add_bias))
             
             # weight initialization
             self.wx[i].weight = torch.nn.Parameter(torch.Tensor(self.dnn_lay[i],current_input).uniform_(-np.sqrt(0.01/(current_input+self.dnn_lay[i])),np.sqrt(0.01/(current_input+self.dnn_lay[i]))))
             self.wx[i].bias = torch.nn.Parameter(torch.zeros(self.dnn_lay[i]))
             
             current_input=self.dnn_lay[i]
             
        self.out_dim=current_input
         
    def forward(self, x):
        
      # Applying Layer/Batch Norm
      if bool(self.dnn_use_laynorm_inp):
        x=self.ln0((x))
        
      if bool(self.dnn_use_batchnorm_inp):
        x=self.bn0((x))
        
      for i in range(self.N_dnn_lay):
           
          if self.dnn_use_laynorm[i] and not(self.dnn_use_batchnorm[i]):
           x = self.drop[i](self.act[i](self.ln[i](self.wx[i](x))))
          
          if self.dnn_use_batchnorm[i] and not(self.dnn_use_laynorm[i]):
           x = self.drop[i](self.act[i](self.bn[i](self.wx[i](x))))
           
          if self.dnn_use_batchnorm[i]==True and self.dnn_use_laynorm[i]==True:
           x = self.drop[i](self.act[i](self.bn[i](self.ln[i](self.wx[i](x)))))
          
          if self.dnn_use_batchnorm[i]==False and self.dnn_use_laynorm[i]==False:
           x = self.drop[i](self.act[i](self.wx[i](x)))
            
          
      return x
########################################################## # pytorch-kaldi v.0.1 # Author: <NAME>, <NAME> # Mila, University of Montreal # October 2018 ########################################################## # input layer normalization # input batch normalization # Initialization of hidden layers # dropout # activation # layer norm initialization # Linear operations # weight initialization # Applying Layer/Batch Norm