EpHod / model /ephod /training /nn_models.py

Upload folder using huggingface_hub

bae5726 verified about 2 months ago

22.9 kB

"""
Neural network models
"""


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F




def torchActivation(activation='elu'):
    '''Return an activation function from torch.nn'''

    if activation == 'relu':
        return nn.ReLU()
    elif activation == 'leaky_relu':
        return nn.LeakyReLU()
    elif activation == 'elu':
        return nn.ELU()
    elif activation == 'selu':
        return nn.SELU()
    elif activation == 'gelu':
        return nn.GELU()




def count_parameters(model):
    '''Return a count of parameters and tensor shape of PyTorch model''' 
    
    counted = {}
    total = 0
    for name, param in model.named_parameters():
        if param.requires_grad:
            count = param.numel()
            total += count
            counted[name] = count
    counted['FULL_MODEL'] = total

    return counted




class ResidualDense(nn.Module):
    '''A single dense layer with residual connection'''
    
    def __init__(self, dim=2560, dropout=0.1, activation='elu', random_seed=0):
        
        super(ResidualDense, self).__init__()
        _ = torch.manual_seed(random_seed)
        self.dense = nn.Linear(dim, dim)
        self.batchnorm = nn.BatchNorm1d(dim)
        self.activation = torchActivation(activation)
        self.dropout = nn.Dropout(dropout)


    def forward(self, x):

        x0 = x
        x = self.dense(x)
        x = self.batchnorm(x)
        x = self.activation(x)        
        x = self.dropout(x)
        x = x0 + x
        
        return x




class LightAttention(nn.Module):
    '''Convolution model with attention to learn pooled representations from embeddings'''

    def __init__(self, dim=1280, kernel_size=7, random_seed=0):
        
        super(LightAttention, self).__init__()
        _ = torch.manual_seed(random_seed)        
        samepad = kernel_size // 2
        self.values_conv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=samepad)
        self.weights_conv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=samepad)
        self.softmax = nn.Softmax(dim=-1)
    
    
    def forward(self, x, mask=None):
        if mask is None:
            mask = torch.ones(x.shape[0], x.shape[2], dtype=torch.int32)  # Don't mask out
        values = self.values_conv(x)
        values = values.masked_fill(mask[:,None,:]==0, -1e6)
        weights = self.weights_conv(x)
        weights = weights.masked_fill(mask[:,None,:]==0, -1e6)
        weights = self.softmax(weights)
        x_sum = torch.sum(values * weights, dim=-1) # Attention-weighted pooling
        x_max, _ = torch.max(values, dim=-1) # Max pooling
        x = torch.cat([x_sum, x_max], dim=1)
        
        return x, weights




class ResidualLightAttention(nn.Module):
    '''Model consisting of light attention followed by residual dense layers'''
    
    def __init__(self, dim=1280, kernel_size=9, dropout=0.5,
                 activation='relu', res_blocks=4, random_seed=0):

        super(ResidualLightAttention, self).__init__()
        torch.manual_seed(random_seed)
        self.light_attention = LightAttention(dim, kernel_size, random_seed)
        self.batchnorm = nn.BatchNorm1d(2 * dim)                
        self.dropout = nn.Dropout(dropout)        
        self.residual_dense = nn.ModuleList()        
        for i in range(res_blocks):
            self.residual_dense.append(
                ResidualDense(2 * dim, dropout, activation, random_seed)
                )
        self.output = nn.Linear(2 * dim, 1)
        
        
    def forward(self, x, mask=None):

        x, weights = self.light_attention(x, mask)
        x = self.batchnorm(x)
        x = self.dropout(x)
        for layer in self.residual_dense:
            x = layer(x)
        y = self.output(x).flatten()
        
        return [y, x, weights]




class FeedForwardNetwork(nn.Module):
    '''A feed-forward neural network (fully-connected)'''
    
    def __init__(self, input_dim=1280, hidden_dim=128, num_layers=1, dropout=0.25, activation='relu', 
                 residual=False):
        
        super().__init__()
        self.num_layers = num_layers
        self.residual = residual

        # First dense layer
        # Separate first dense layer from the other dense layers to allow residual connection
        self.first_dense = nn.ModuleList(
            [
                nn.Linear(input_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),                                         
                torchActivation(activation),
                nn.Dropout(dropout) if activation != 'selu' else nn.AlphaDropout(dropout)
            ]
        )
        
        # Other dense layers after first dense and residual connection
        if num_layers > 1:
            self.other_dense = nn.ModuleList()
            for i in range(1, num_layers):
                self.other_dense.extend(
                    [
                        nn.Linear(hidden_dim, hidden_dim),
                        nn.BatchNorm1d(hidden_dim),
                        torchActivation(activation),
                        nn.Dropout(dropout) if activation != 'selu' else nn.AlphaDropout(dropout)
                    ]
                )
        # Output layer
        self.output_layer = nn.Linear(hidden_dim, 1)


    def forward(self, x):

        # First dense layer
        for layer in self.first_dense:
            x = layer(x) 
            
        # Other dense layers
        if self.num_layers > 1:
            x0 = x
            for layer in self.other_dense:
                x = layer(x)  
            # Residual connection
            if self.residual:
                x = x0 + x 
        # Output layer
        y = self.output_layer(x).flatten()
        
        return y




class CNN(nn.Module):

    def __init__(self, input_channel=20, input_length=1024, start_conv_channel=32, kernel_size=3,
                 num_conv_layers=8, conv_dropout=0.25, pooltype='max', dense_dim=128, num_dense_layers=1,
                 dense_dropout=0.25, activation='relu'):
        
        super().__init__()
        # Conv layers
        self.conv_layers = nn.ModuleList() 
        for i in range(num_conv_layers):
            in_channel = input_channel if i==0 else out_channel
            out_channel = start_conv_channel * (2 ** i)
            self.conv_layers.append(nn.Conv1d(in_channel, out_channel, kernel_size=kernel_size, 
                                              padding=(kernel_size // 2)))
            self.conv_layers.append(nn.BatchNorm1d(out_channel))
            self.conv_layers.append(torchActivation(activation))
            if pooltype=='max' and i>0:
                self.conv_layers.append(nn.MaxPool1d(kernel_size=2, stride=2))
            elif pooltype=='average' and i>0:
                self.conv_layers.append(nn.AvgPool1d(kernel_size=2, stride=2))
            self.conv_layers.append(nn.Dropout(conv_dropout))
        # FNN layers
        self.flatsize = start_conv_channel * input_length
        self.fnn = FeedForwardNetwork(input_dim=self.flatsize, hidden_dim=dense_dim, num_layers=num_dense_layers, 
                                      dropout=dense_dropout, activation=activation, residual=True)

    def forward(self, x):

        for layer in self.conv_layers:
            x = layer(x)
        x = x.view(-1, self.flatsize)
        y = self.fnn(x)
        
        return y




class CausalConv1d(nn.Module):
    '''Causal convolution'''
    
    def __init__(self, in_channels, out_channels, kernel_size, dilation=1):

        super().__init__()

        self.pad = int(np.ceil(dilation * (kernel_size - 1) / 2)) # Padding to ensure causal convolution
        self.conv_layer = nn.Conv1d(in_channels, out_channels, padding=0,kernel_size=kernel_size, 
                                    dilation=dilation)
    
    def forward(self, x):
        
        len1 = x.shape[-1]
        x = F.pad(x, (self.pad * 2, 0))
        x = self.conv_layer(x)
        len2 = x.shape[-1]
        start = len2 - len1
        x = x[:,:,start:len2]

        return x




class DCNNBlock(nn.Module):

    def __init__(self,channel_size=32, num_layers=6,dropout=0.5, activation='elu'):
        
        super().__init__()
        dilation_rates = (2. ** np.arange(num_layers)).astype(np.int32)
        self.layers = nn.ModuleList()
        for d in dilation_rates:
            self.layers.append(CausalConv1d(channel_size, channel_size, kernel_size=2, dilation=d))
            self.layers.append(nn.BatchNorm1d(channel_size))
            self.layers.append(torchActivation(activation))
        self.layers.append(nn.Dropout(dropout))
        
    def forward(self, x):
        
        x0 = x
        for layer in self.layers:
            x = layer(x)
        x = x0 + x

        return x




class DCNN(nn.Module):

    def __init__(self, input_channel=20, conv_channel=48, num_blocks=6, layers_per_block=10,
                 block_dropout=0.5, dense_dim=128, num_dense_layers=2, dense_dropout=0.25,
                 activation='elu'):
        
        super().__init__()
        self.embed_layer = nn.Conv1d(input_channel, conv_channel, kernel_size=1, bias=False)
        
        # DCNN blocks
        self.dcnn_blocks = nn.ModuleList()    
        for block in range(num_blocks):
            self.dcnn_blocks.append(DCNNBlock(channel_size=conv_channel, num_layers=layers_per_block,
                                              dropout=block_dropout, activation=activation))
        
        self.fnn = FeedForwardNetwork(input_dim=conv_channel, hidden_dim=dense_dim, 
                                      num_layers=num_dense_layers, dropout=dense_dropout, 
                                      activation=activation, residual=True)
    
    def forward(self, x, mask=None):

        x = self.embed_layer(x)
        # Dilated convolutions
        for layer in self.dcnn_blocks:
            x = layer(x)
        if mask is None:
            mask = torch.ones(x.shape[0], x.shape[2], dtype=torch.int32)
        # Average pooling
        x = x.masked_fill(mask[:,None,:] == 0, 0)
        x = torch.divide(torch.sum(x, dim=-1), torch.sum(mask, dim=-1)[:,None])
        
        # FNN
        y = self.fnn(x).flatten()
        
        return y




class RNN(nn.Module):

    def __init__(self, input_channel=20, input_length=1024, gru_dim=1024, conv_downsample=1, 
                 conv_dropout=0.25, dense_dim=128, dense_dropout=0.25, activation='relu'):
        
        super().__init__()
        
        # Convolutional layers to downsample
        self.conv_downsample = conv_downsample
        self.conv_layers = nn.ModuleList() 
        pad = (conv_downsample // 2 - 1)
        pad = max(pad, 0)
        self.conv_layers.append(nn.Conv1d(input_channel, gru_dim, kernel_size=conv_downsample, 
                                          stride=conv_downsample, padding=pad))
        self.conv_layers.append(nn.BatchNorm1d(gru_dim))
        self.conv_layers.append(torchActivation(activation))
        self.conv_layers.append(nn.Dropout(conv_dropout))
        
        # RNN layer
        self.gru_layer = nn.GRU(gru_dim, gru_dim, batch_first=True, bidirectional=False)
       
        # Dense layer
        self.dense_layers = nn.ModuleList()
        self.dense_layers.extend(
            [
                nn.Linear(gru_dim, dense_dim),
                nn.BatchNorm1d(dense_dim),                    
                torchActivation(activation),
                nn.Dropout(dense_dropout) if activation != 'selu' else nn.AlphaDropout(dense_dropout)
            ]
        )

        # Output layer
        self.output_layer = nn.Linear(dense_dim, 1) 



    def forward(self, x, mask=None):
        
        if mask is None:
            mask = torch.tensor(np.ones((x.shape[0], x.shape[-1])), dtype=torch.float32)

        # Downsample with convolution
        for layer in self.conv_layers:
            x = layer(x)

        # RNN
        x = x.transpose(2, 1) # Transpose from [batch, feat, seqlen] to [batch, sequence, features]
        x = self.gru_layer(x)[0]
        lastind = torch.sum(mask, axis=1, dtype=int) # index of last residue in sequence
        lastind = (lastind // self.conv_downsample) - 1  # index after conv. downsampling
        x = [x[i][lastind[i]] for i in range(len(lastind))] # get hidden state of last residue for each sequence
        x = torch.stack(x)
        
        # Dense layer
        for layer in self.dense_layers:
            x = layer(x)
        
        # Output layer
        x = self.output_layer(x).flatten()

        return x




class AttentionTopModel(nn.Module):
    '''
    Class to implement a top model based on LAT architecture to predict a property from 
    embeddings.
    
    Parameters
    -----------
    input_dim : int
        Channel dimension of input embedding, 1280 for ESM-1b        
    attention_dim : int
        Channel dimension of transformed values and attention weights
    kernel_size : int
        Kernel size of convolution 
    conv_dropout : float
        Dropout applied to values embedding after convolution 
    perceptive : bool
        If True, concatenate input embedding with average pooled embedding to derive 
        a perceptive embedding.        
    dense_dim : int
        Number of units in hidden dense layer
    dense_dropout : float
        Dropout of dense hidden layers.
    activation : str
        Activation function of dense hidden layer.
    random_seed : int
        Random seed for reproducibillity
     
     '''


    def __init__(self, 
                 input_dim=1280,
                 attention_dim=1280,
                 kernel_size=9,
                 conv_dropout=0.25,
                 perceptive=True,                 
                 dense_dim=512,
                 dense_dropout=0.25,
                 num_layers=1, 
                 residual=False,
                 activation='relu',
                 random_seed=0):
        

        super(AttentionTopModel, self).__init__()


        # Initialize class parameters  
        self.input_dim = input_dim            
        self.attention_dim = attention_dim   
        self.kernel_size = kernel_size       
        self.conv_dropout = conv_dropout    
        self.perceptive = perceptive                 
        self.dense_dim = dense_dim           
        self.dense_dropout = dense_dropout   
        self.num_layers = num_layers
        self.residual = residual
        self.activation = activation         
        self.random_seed = random_seed       
        _ = torch.manual_seed(self.random_seed)

            
        # A convolution layer to learn values from embeddings
        self.values_conv = nn.Conv1d(
            2 * self.input_dim if perceptive else self.input_dim,
            self.attention_dim, 
            kernel_size=self.kernel_size,
            stride=1,
            padding=(self.kernel_size // 2)
            )
        self.values_dropout = nn.Dropout(self.conv_dropout) 
        
        
        # A convolution and softmax layer to learn attention weights that sum to 1
        self.attn_conv = nn.Conv1d(
            2 * self.input_dim if perceptive else self.input_dim,
            self.attention_dim, 
            kernel_size=self.kernel_size,
            stride=1,
            padding=(self.kernel_size // 2)
            )
        self.attn_softmax = nn.Softmax(dim=-1)
        
        # Dense(residual) layers on top of the attention-learned embeddings
        self.dense_layers = FeedForwardNetwork(
            input_dim=(2 * self.attention_dim), 
            hidden_dims=[self.dense_dim] * self.num_layers, 
            dropout=self.dense_dropout, 
            activation=self.activation, 
            residual=self.residual, 
            random_seed=self.random_seed
            )
                                                     

    def forward(self, x, mask=None):
        '''
        Forward pass of network. 
        Parameters
        ------------
        x : torch.Tensor
            Shape is [batch_size, embeddings_dim, seq_len]
        mask : torch.Tensor or None
            Boolean tensor indicating padded (0) and non-padded (1) positions. Shape is 
            [batch_size, seq_len]. If mask is None, a tensor of ones is used, 
            indicating that padded positions are not masked out in the forward 
            computations.
        Returns
        -------
        output: list
            A list of predicted values (shape is [batch_size, 1]), embedding of the
            hidden layer before the output layer (shape is [batch_size, dense_dim]), and
            the masked attention weights (shape is [batch_size, attention_dim, seq_len])
        '''
        
        # Masking for  padded positions
        if mask is None:
            # If mask is None, use all positions including padded positions
            mask = torch.ones(x.shape[0], x.shape[2], dtype=torch.int32)
        
        # Derive a perceptive tensor by concatenating mean pooled embeddings to the input 
        # Output shape of perceptive tensor is [batch_size, 2 * input_dim, seq_length]
        if self.perceptive:
            # Mask out padded positions
            x_perc = x.masked_fill(mask[:,None,:] == 0, 0)  
            # Average pooling over unpadded positions
            x_perc = torch.divide(torch.sum(x_perc, dim=-1), 
                                  torch.sum(mask, dim=1)[:,None]) 
            # Repeat pooled tensor to match sequence length
            x_perc = x_perc[:,:,None].repeat(1,1,x.shape[-1]) 
            # Concatenate input embeddings and perceptive tensor along feature dimension
            x = torch.cat((x, x_perc), axis=1) 
        
        # Derive values and attention weights from input embeddings with convolution
        # Output shapes for values/weights is [batch_size, attention_dim, seq_length]
        values = self.values_conv(x)
        values = self.values_dropout(values)
        values = values.masked_fill(mask[:,None,:] == 0, -1e6)  # Mask out for max pooling

        # Attention weights that sum to 1
        attn = self.attn_conv(x)   
        attn = attn.masked_fill(mask[:,None,:] == 0, -1e6) # Mask out for average pooling
        attn = self.attn_softmax(attn)
        
        # Derive attention-learned embeddings by multiplying values and attention weights
        # Output shape is [batch_size, 2 * attention_dim]
        z_sum = torch.sum(values * attn, dim=-1) 
        z_max, _ = torch.max(values, dim=-1)
        z = torch.cat([z_sum, z_max], dim=1)

        # Dense layers
        y, z = self.dense_layers(z)

        
        return [y, z, attn]




class DilatedConvTopModel(nn.Module):

    def __init__(self, 
                 input_dim=1280,
                 input_len=1024,
                 attention_kernel_size=9,                 
                 conv_dim=64,
                 conv_kernel_size=3,
                 conv_dropout=0.1,
                 dilated_layers=9,                 
                 dense_dim=512,
                 dense_layers=1,                  
                 dense_dropout=0.25,
                 activation='relu',
                 random_seed=0):
        
        super(DilatedConvTopModel, self).__init__()
        _ = torch.manual_seed(random_seed)
        
        # A convolution layer to learn transformed values from embeddings
        self.values_conv_layer = nn.Conv1d(
            input_dim,
            conv_dim, 
            kernel_size=attention_kernel_size,
            stride=1, 
            dilation=1, 
            padding=(attention_kernel_size // 2)
            )        
        
        
        # A convolution and softmax layer to learn attention weights that sum to 1
        # Attention-weighted values are obtained by element-wise multiplication with weights
        self.attention_conv_layer = nn.Conv1d(
            input_dim,
            conv_dim, 
            kernel_size=attention_kernel_size,
            stride=1, 
            dilation=1, 
            padding=(attention_kernel_size // 2)
            )       
        self.attention_softmax = nn.Softmax(dim=-1)
        self.conv_dropout_layer = nn.Dropout(conv_dropout) if activation != 'selu' \
                    else nn.AlphaDropout(conv_dropout)
        
        
        # Convolutional layers with increasing dilation rates
        self.dilated_conv_layers = nn.ModuleList()
        dilations = np.power(2, np.arange(dilated_layers)) 

        for d in dilations:

            same_padding = self.get_same_padding_size(conv_kernel_size, d)
            self.dilated_conv_layers.append(nn.Conv1d(conv_dim, 
                                                      conv_dim,
                                                      kernel_size=conv_kernel_size,
                                                      padding=same_padding,
                                                      dilation=d)
                              )
            self.dilated_conv_layers.append(nn.BatchNorm1d(conv_dim))
            self.dilated_conv_layers.append(torchActivation(activation))
            
        
        # Dense layers
        self.dense_layers = FeedForwardNetwork(input_dim=(input_len * conv_dim), 
                                               hidden_dims=[dense_dim] * dense_layers, 
                                               dropout=dense_dropout, 
                                               activation=activation, 
                                               residual=(dense_layers > 1), 
                                               random_seed=random_seed)

    
    def get_same_padding_size(self, k, d):

        padsize = int(
            np.ceil(
                d * (k - 1) / 2
                )
            )
        
        return padsize
    
    
    def forward(self, x, mask=None):
        
        # Masking for padded positions
        if mask is None:
            # If mask is None, use all positions including padded positions
            mask = torch.ones(x.shape[0], x.shape[2], dtype=torch.int32)
        
        # Derive attention-weighted values from input embeddings
        # Shapes for values/weights is [batch_size, attention_dim, seq_length]
        x_values = self.values_conv_layer(x)
        x_weights = self.attention_conv_layer(x)   
        x_weights = x_weights.masked_fill(mask[:,None,:] == 0, -1e6) 
        x_weights = self.attention_softmax(x_weights) # Weights sum to 1
        x = (x_values * x_weights)  
        x = self.conv_dropout_layer(x)
        
        # Dilated convolution layers
        x0 = x
        for layer in self.dilated_conv_layers:
            x = layer(x)
        x = x0 + x  # Residual connection
        
        # Dense layers
        x = x.flatten(start_dim=1, end_dim=2)
        y, z = self.dense_layers(x)
        
        return [y, z, x_weights]     

	"""
	Neural network models
	"""


	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F




	def torchActivation(activation='elu'):
	'''Return an activation function from torch.nn'''

	if activation == 'relu':
	return nn.ReLU()
	elif activation == 'leaky_relu':
	return nn.LeakyReLU()
	elif activation == 'elu':
	return nn.ELU()
	elif activation == 'selu':
	return nn.SELU()
	elif activation == 'gelu':
	return nn.GELU()




	def count_parameters(model):
	'''Return a count of parameters and tensor shape of PyTorch model'''

	counted = {}
	total = 0
	for name, param in model.named_parameters():
	if param.requires_grad:
	count = param.numel()
	total += count
	counted[name] = count
	counted['FULL_MODEL'] = total

	return counted




	class ResidualDense(nn.Module):
	'''A single dense layer with residual connection'''

	def __init__(self, dim=2560, dropout=0.1, activation='elu', random_seed=0):

	super(ResidualDense, self).__init__()
	_ = torch.manual_seed(random_seed)
	self.dense = nn.Linear(dim, dim)
	self.batchnorm = nn.BatchNorm1d(dim)
	self.activation = torchActivation(activation)
	self.dropout = nn.Dropout(dropout)


	def forward(self, x):

	x0 = x
	x = self.dense(x)
	x = self.batchnorm(x)
	x = self.activation(x)
	x = self.dropout(x)
	x = x0 + x

	return x




	class LightAttention(nn.Module):
	'''Convolution model with attention to learn pooled representations from embeddings'''

	def __init__(self, dim=1280, kernel_size=7, random_seed=0):

	super(LightAttention, self).__init__()
	_ = torch.manual_seed(random_seed)
	samepad = kernel_size // 2
	self.values_conv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=samepad)
	self.weights_conv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=samepad)
	self.softmax = nn.Softmax(dim=-1)


	def forward(self, x, mask=None):
	if mask is None:
	mask = torch.ones(x.shape[0], x.shape[2], dtype=torch.int32) # Don't mask out
	values = self.values_conv(x)
	values = values.masked_fill(mask[:,None,:]==0, -1e6)
	weights = self.weights_conv(x)
	weights = weights.masked_fill(mask[:,None,:]==0, -1e6)
	weights = self.softmax(weights)
	x_sum = torch.sum(values * weights, dim=-1) # Attention-weighted pooling
	x_max, _ = torch.max(values, dim=-1) # Max pooling
	x = torch.cat([x_sum, x_max], dim=1)

	return x, weights




	class ResidualLightAttention(nn.Module):
	'''Model consisting of light attention followed by residual dense layers'''

	def __init__(self, dim=1280, kernel_size=9, dropout=0.5,
	activation='relu', res_blocks=4, random_seed=0):

	super(ResidualLightAttention, self).__init__()
	torch.manual_seed(random_seed)
	self.light_attention = LightAttention(dim, kernel_size, random_seed)
	self.batchnorm = nn.BatchNorm1d(2 * dim)
	self.dropout = nn.Dropout(dropout)
	self.residual_dense = nn.ModuleList()
	for i in range(res_blocks):
	self.residual_dense.append(
	ResidualDense(2 * dim, dropout, activation, random_seed)
	)
	self.output = nn.Linear(2 * dim, 1)


	def forward(self, x, mask=None):

	x, weights = self.light_attention(x, mask)
	x = self.batchnorm(x)
	x = self.dropout(x)
	for layer in self.residual_dense:
	x = layer(x)
	y = self.output(x).flatten()

	return [y, x, weights]




	class FeedForwardNetwork(nn.Module):
	'''A feed-forward neural network (fully-connected)'''

	def __init__(self, input_dim=1280, hidden_dim=128, num_layers=1, dropout=0.25, activation='relu',
	residual=False):

	super().__init__()
	self.num_layers = num_layers
	self.residual = residual

	# First dense layer
	# Separate first dense layer from the other dense layers to allow residual connection
	self.first_dense = nn.ModuleList(
	[
	nn.Linear(input_dim, hidden_dim),
	nn.BatchNorm1d(hidden_dim),
	torchActivation(activation),
	nn.Dropout(dropout) if activation != 'selu' else nn.AlphaDropout(dropout)
	]
	)

	# Other dense layers after first dense and residual connection
	if num_layers > 1:
	self.other_dense = nn.ModuleList()
	for i in range(1, num_layers):
	self.other_dense.extend(
	[
	nn.Linear(hidden_dim, hidden_dim),
	nn.BatchNorm1d(hidden_dim),
	torchActivation(activation),
	nn.Dropout(dropout) if activation != 'selu' else nn.AlphaDropout(dropout)
	]
	)
	# Output layer
	self.output_layer = nn.Linear(hidden_dim, 1)


	def forward(self, x):

	# First dense layer
	for layer in self.first_dense:
	x = layer(x)

	# Other dense layers
	if self.num_layers > 1:
	x0 = x
	for layer in self.other_dense:
	x = layer(x)
	# Residual connection
	if self.residual:
	x = x0 + x
	# Output layer
	y = self.output_layer(x).flatten()

	return y




	class CNN(nn.Module):

	def __init__(self, input_channel=20, input_length=1024, start_conv_channel=32, kernel_size=3,
	num_conv_layers=8, conv_dropout=0.25, pooltype='max', dense_dim=128, num_dense_layers=1,
	dense_dropout=0.25, activation='relu'):

	super().__init__()
	# Conv layers
	self.conv_layers = nn.ModuleList()
	for i in range(num_conv_layers):
	in_channel = input_channel if i==0 else out_channel
	out_channel = start_conv_channel * (2 ** i)
	self.conv_layers.append(nn.Conv1d(in_channel, out_channel, kernel_size=kernel_size,
	padding=(kernel_size // 2)))
	self.conv_layers.append(nn.BatchNorm1d(out_channel))
	self.conv_layers.append(torchActivation(activation))
	if pooltype=='max' and i>0:
	self.conv_layers.append(nn.MaxPool1d(kernel_size=2, stride=2))
	elif pooltype=='average' and i>0:
	self.conv_layers.append(nn.AvgPool1d(kernel_size=2, stride=2))
	self.conv_layers.append(nn.Dropout(conv_dropout))
	# FNN layers
	self.flatsize = start_conv_channel * input_length
	self.fnn = FeedForwardNetwork(input_dim=self.flatsize, hidden_dim=dense_dim, num_layers=num_dense_layers,
	dropout=dense_dropout, activation=activation, residual=True)

	def forward(self, x):

	for layer in self.conv_layers:
	x = layer(x)
	x = x.view(-1, self.flatsize)
	y = self.fnn(x)

	return y




	class CausalConv1d(nn.Module):
	'''Causal convolution'''

	def __init__(self, in_channels, out_channels, kernel_size, dilation=1):

	super().__init__()

	self.pad = int(np.ceil(dilation * (kernel_size - 1) / 2)) # Padding to ensure causal convolution
	self.conv_layer = nn.Conv1d(in_channels, out_channels, padding=0,kernel_size=kernel_size,
	dilation=dilation)

	def forward(self, x):

	len1 = x.shape[-1]
	x = F.pad(x, (self.pad * 2, 0))
	x = self.conv_layer(x)
	len2 = x.shape[-1]
	start = len2 - len1
	x = x[:,:,start:len2]

	return x




	class DCNNBlock(nn.Module):

	def __init__(self,channel_size=32, num_layers=6,dropout=0.5, activation='elu'):

	super().__init__()
	dilation_rates = (2. ** np.arange(num_layers)).astype(np.int32)
	self.layers = nn.ModuleList()
	for d in dilation_rates:
	self.layers.append(CausalConv1d(channel_size, channel_size, kernel_size=2, dilation=d))
	self.layers.append(nn.BatchNorm1d(channel_size))
	self.layers.append(torchActivation(activation))
	self.layers.append(nn.Dropout(dropout))

	def forward(self, x):

	x0 = x
	for layer in self.layers:
	x = layer(x)
	x = x0 + x

	return x




	class DCNN(nn.Module):

	def __init__(self, input_channel=20, conv_channel=48, num_blocks=6, layers_per_block=10,
	block_dropout=0.5, dense_dim=128, num_dense_layers=2, dense_dropout=0.25,
	activation='elu'):

	super().__init__()
	self.embed_layer = nn.Conv1d(input_channel, conv_channel, kernel_size=1, bias=False)

	# DCNN blocks
	self.dcnn_blocks = nn.ModuleList()
	for block in range(num_blocks):
	self.dcnn_blocks.append(DCNNBlock(channel_size=conv_channel, num_layers=layers_per_block,
	dropout=block_dropout, activation=activation))

	self.fnn = FeedForwardNetwork(input_dim=conv_channel, hidden_dim=dense_dim,
	num_layers=num_dense_layers, dropout=dense_dropout,
	activation=activation, residual=True)

	def forward(self, x, mask=None):

	x = self.embed_layer(x)
	# Dilated convolutions
	for layer in self.dcnn_blocks:
	x = layer(x)
	if mask is None:
	mask = torch.ones(x.shape[0], x.shape[2], dtype=torch.int32)
	# Average pooling
	x = x.masked_fill(mask[:,None,:] == 0, 0)
	x = torch.divide(torch.sum(x, dim=-1), torch.sum(mask, dim=-1)[:,None])

	# FNN
	y = self.fnn(x).flatten()

	return y




	class RNN(nn.Module):

	def __init__(self, input_channel=20, input_length=1024, gru_dim=1024, conv_downsample=1,
	conv_dropout=0.25, dense_dim=128, dense_dropout=0.25, activation='relu'):

	super().__init__()

	# Convolutional layers to downsample
	self.conv_downsample = conv_downsample
	self.conv_layers = nn.ModuleList()
	pad = (conv_downsample // 2 - 1)
	pad = max(pad, 0)
	self.conv_layers.append(nn.Conv1d(input_channel, gru_dim, kernel_size=conv_downsample,
	stride=conv_downsample, padding=pad))
	self.conv_layers.append(nn.BatchNorm1d(gru_dim))
	self.conv_layers.append(torchActivation(activation))
	self.conv_layers.append(nn.Dropout(conv_dropout))

	# RNN layer
	self.gru_layer = nn.GRU(gru_dim, gru_dim, batch_first=True, bidirectional=False)

	# Dense layer
	self.dense_layers = nn.ModuleList()
	self.dense_layers.extend(
	[
	nn.Linear(gru_dim, dense_dim),
	nn.BatchNorm1d(dense_dim),
	torchActivation(activation),
	nn.Dropout(dense_dropout) if activation != 'selu' else nn.AlphaDropout(dense_dropout)
	]
	)

	# Output layer
	self.output_layer = nn.Linear(dense_dim, 1)



	def forward(self, x, mask=None):

	if mask is None:
	mask = torch.tensor(np.ones((x.shape[0], x.shape[-1])), dtype=torch.float32)

	# Downsample with convolution
	for layer in self.conv_layers:
	x = layer(x)

	# RNN
	x = x.transpose(2, 1) # Transpose from [batch, feat, seqlen] to [batch, sequence, features]
	x = self.gru_layer(x)[0]
	lastind = torch.sum(mask, axis=1, dtype=int) # index of last residue in sequence
	lastind = (lastind // self.conv_downsample) - 1 # index after conv. downsampling
	x = [x[i][lastind[i]] for i in range(len(lastind))] # get hidden state of last residue for each sequence
	x = torch.stack(x)

	# Dense layer
	for layer in self.dense_layers:
	x = layer(x)

	# Output layer
	x = self.output_layer(x).flatten()

	return x




	class AttentionTopModel(nn.Module):
	'''
	Class to implement a top model based on LAT architecture to predict a property from
	embeddings.

	Parameters
	-----------
	input_dim : int
	Channel dimension of input embedding, 1280 for ESM-1b
	attention_dim : int
	Channel dimension of transformed values and attention weights
	kernel_size : int
	Kernel size of convolution
	conv_dropout : float
	Dropout applied to values embedding after convolution
	perceptive : bool
	If True, concatenate input embedding with average pooled embedding to derive
	a perceptive embedding.
	dense_dim : int
	Number of units in hidden dense layer
	dense_dropout : float
	Dropout of dense hidden layers.
	activation : str
	Activation function of dense hidden layer.
	random_seed : int
	Random seed for reproducibillity

	'''


	def __init__(self,
	input_dim=1280,
	attention_dim=1280,
	kernel_size=9,
	conv_dropout=0.25,
	perceptive=True,
	dense_dim=512,
	dense_dropout=0.25,
	num_layers=1,
	residual=False,
	activation='relu',
	random_seed=0):


	super(AttentionTopModel, self).__init__()


	# Initialize class parameters
	self.input_dim = input_dim
	self.attention_dim = attention_dim
	self.kernel_size = kernel_size
	self.conv_dropout = conv_dropout
	self.perceptive = perceptive
	self.dense_dim = dense_dim
	self.dense_dropout = dense_dropout
	self.num_layers = num_layers
	self.residual = residual
	self.activation = activation
	self.random_seed = random_seed
	_ = torch.manual_seed(self.random_seed)


	# A convolution layer to learn values from embeddings
	self.values_conv = nn.Conv1d(
	2 * self.input_dim if perceptive else self.input_dim,
	self.attention_dim,
	kernel_size=self.kernel_size,
	stride=1,
	padding=(self.kernel_size // 2)
	)
	self.values_dropout = nn.Dropout(self.conv_dropout)


	# A convolution and softmax layer to learn attention weights that sum to 1
	self.attn_conv = nn.Conv1d(
	2 * self.input_dim if perceptive else self.input_dim,
	self.attention_dim,
	kernel_size=self.kernel_size,
	stride=1,
	padding=(self.kernel_size // 2)
	)
	self.attn_softmax = nn.Softmax(dim=-1)

	# Dense(residual) layers on top of the attention-learned embeddings
	self.dense_layers = FeedForwardNetwork(
	input_dim=(2 * self.attention_dim),
	hidden_dims=[self.dense_dim] * self.num_layers,
	dropout=self.dense_dropout,
	activation=self.activation,
	residual=self.residual,
	random_seed=self.random_seed
	)


	def forward(self, x, mask=None):
	'''
	Forward pass of network.

	Parameters
	------------
	x : torch.Tensor
	Shape is [batch_size, embeddings_dim, seq_len]
	mask : torch.Tensor or None
	Boolean tensor indicating padded (0) and non-padded (1) positions. Shape is
	[batch_size, seq_len]. If mask is None, a tensor of ones is used,
	indicating that padded positions are not masked out in the forward
	computations.

	Returns
	-------
	output: list
	A list of predicted values (shape is [batch_size, 1]), embedding of the
	hidden layer before the output layer (shape is [batch_size, dense_dim]), and
	the masked attention weights (shape is [batch_size, attention_dim, seq_len])
	'''

	# Masking for padded positions
	if mask is None:
	# If mask is None, use all positions including padded positions
	mask = torch.ones(x.shape[0], x.shape[2], dtype=torch.int32)

	# Derive a perceptive tensor by concatenating mean pooled embeddings to the input
	# Output shape of perceptive tensor is [batch_size, 2 * input_dim, seq_length]
	if self.perceptive:
	# Mask out padded positions
	x_perc = x.masked_fill(mask[:,None,:] == 0, 0)
	# Average pooling over unpadded positions
	x_perc = torch.divide(torch.sum(x_perc, dim=-1),
	torch.sum(mask, dim=1)[:,None])
	# Repeat pooled tensor to match sequence length
	x_perc = x_perc[:,:,None].repeat(1,1,x.shape[-1])
	# Concatenate input embeddings and perceptive tensor along feature dimension
	x = torch.cat((x, x_perc), axis=1)

	# Derive values and attention weights from input embeddings with convolution
	# Output shapes for values/weights is [batch_size, attention_dim, seq_length]
	values = self.values_conv(x)
	values = self.values_dropout(values)
	values = values.masked_fill(mask[:,None,:] == 0, -1e6) # Mask out for max pooling

	# Attention weights that sum to 1
	attn = self.attn_conv(x)
	attn = attn.masked_fill(mask[:,None,:] == 0, -1e6) # Mask out for average pooling
	attn = self.attn_softmax(attn)

	# Derive attention-learned embeddings by multiplying values and attention weights
	# Output shape is [batch_size, 2 * attention_dim]
	z_sum = torch.sum(values * attn, dim=-1)
	z_max, _ = torch.max(values, dim=-1)
	z = torch.cat([z_sum, z_max], dim=1)

	# Dense layers
	y, z = self.dense_layers(z)


	return [y, z, attn]




	class DilatedConvTopModel(nn.Module):

	def __init__(self,
	input_dim=1280,
	input_len=1024,
	attention_kernel_size=9,
	conv_dim=64,
	conv_kernel_size=3,
	conv_dropout=0.1,
	dilated_layers=9,
	dense_dim=512,
	dense_layers=1,
	dense_dropout=0.25,
	activation='relu',
	random_seed=0):

	super(DilatedConvTopModel, self).__init__()
	_ = torch.manual_seed(random_seed)

	# A convolution layer to learn transformed values from embeddings
	self.values_conv_layer = nn.Conv1d(
	input_dim,
	conv_dim,
	kernel_size=attention_kernel_size,
	stride=1,
	dilation=1,
	padding=(attention_kernel_size // 2)
	)


	# A convolution and softmax layer to learn attention weights that sum to 1
	# Attention-weighted values are obtained by element-wise multiplication with weights
	self.attention_conv_layer = nn.Conv1d(
	input_dim,
	conv_dim,
	kernel_size=attention_kernel_size,
	stride=1,
	dilation=1,
	padding=(attention_kernel_size // 2)
	)
	self.attention_softmax = nn.Softmax(dim=-1)
	self.conv_dropout_layer = nn.Dropout(conv_dropout) if activation != 'selu' \
	else nn.AlphaDropout(conv_dropout)


	# Convolutional layers with increasing dilation rates
	self.dilated_conv_layers = nn.ModuleList()
	dilations = np.power(2, np.arange(dilated_layers))

	for d in dilations:

	same_padding = self.get_same_padding_size(conv_kernel_size, d)
	self.dilated_conv_layers.append(nn.Conv1d(conv_dim,
	conv_dim,
	kernel_size=conv_kernel_size,
	padding=same_padding,
	dilation=d)
	)
	self.dilated_conv_layers.append(nn.BatchNorm1d(conv_dim))
	self.dilated_conv_layers.append(torchActivation(activation))


	# Dense layers
	self.dense_layers = FeedForwardNetwork(input_dim=(input_len * conv_dim),
	hidden_dims=[dense_dim] * dense_layers,
	dropout=dense_dropout,
	activation=activation,
	residual=(dense_layers > 1),
	random_seed=random_seed)


	def get_same_padding_size(self, k, d):

	padsize = int(
	np.ceil(
	d * (k - 1) / 2
	)
	)

	return padsize


	def forward(self, x, mask=None):

	# Masking for padded positions
	if mask is None:
	# If mask is None, use all positions including padded positions
	mask = torch.ones(x.shape[0], x.shape[2], dtype=torch.int32)

	# Derive attention-weighted values from input embeddings
	# Shapes for values/weights is [batch_size, attention_dim, seq_length]
	x_values = self.values_conv_layer(x)
	x_weights = self.attention_conv_layer(x)
	x_weights = x_weights.masked_fill(mask[:,None,:] == 0, -1e6)
	x_weights = self.attention_softmax(x_weights) # Weights sum to 1
	x = (x_values * x_weights)
	x = self.conv_dropout_layer(x)

	# Dilated convolution layers
	x0 = x
	for layer in self.dilated_conv_layers:
	x = layer(x)
	x = x0 + x # Residual connection

	# Dense layers
	x = x.flatten(start_dim=1, end_dim=2)
	y, z = self.dense_layers(x)

	return [y, z, x_weights]