Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| class DeepConvDTI(nn.Module): | |
| def __init__(self, dropout=0.2, drug_layers=(1024, 512), protein_windows=(10, 15, 20, 25), n_filters=64, | |
| decay=0.0, fc_layers=None, convolution=True, activation=nn.ReLU(), protein_layers=None): | |
| super().__init__() | |
| self.dropout = dropout | |
| self.drug_layers = drug_layers | |
| self.protein_windows = protein_windows | |
| self.filters = n_filters | |
| self.decay = decay | |
| self.fc_layers = fc_layers | |
| self.convolution = convolution | |
| self.activation = activation # Use any nn.Module as the activation function | |
| self.protein_layers = protein_layers | |
| # Define the drug branch of the model | |
| self.drug_branch = [] | |
| for layer_size in drug_layers: | |
| self.drug_branch += [ | |
| nn.LazyLinear(layer_size), | |
| nn.BatchNorm1d(layer_size), | |
| activation, | |
| nn.Dropout(dropout) | |
| ] | |
| self.drug_branch = nn.Sequential(*self.drug_branch) | |
| # Define the protein branch of the model | |
| if convolution: | |
| # Use embedding and convolution layers for protein sequences | |
| self.protein_embedding = nn.Embedding(26, 20) | |
| # Use a list of parallel convolution and pooling layers with different window sizes | |
| self.protein_convs = nn.ModuleList() | |
| for window_size in protein_windows: | |
| conv = nn.Sequential( | |
| nn.Conv1d(20, n_filters, window_size, padding="same"), | |
| nn.BatchNorm1d(n_filters), | |
| activation, | |
| nn.AdaptiveMaxPool1d(1) | |
| ) | |
| self.protein_convs.append(conv) | |
| if protein_layers: | |
| self.protein_branch = [] | |
| for layer_size in protein_layers: | |
| self.protein_branch += [ | |
| nn.LazyLinear(layer_size), | |
| nn.BatchNorm1d(layer_size), | |
| activation, | |
| nn.Dropout(dropout) | |
| ] | |
| self.protein_branch = nn.Sequential(*self.protein_branch) | |
| # Define the final branch of the model that combines the drug and protein branches | |
| self.final_branch = [] | |
| if fc_layers: | |
| # Add additional dense layers for the final branch | |
| for layer_size in fc_layers: | |
| self.final_branch += [ | |
| nn.LazyLinear(layer_size), | |
| nn.BatchNorm1d(layer_size), | |
| activation | |
| ] | |
| self.final_branch = nn.Sequential(*self.final_branch) | |
| def forward(self, input_d, input_p): | |
| # Forward pass of the drug branch | |
| output_d = self.drug_branch(input_d.float()) | |
| # Forward pass of the protein branch | |
| if self.convolution: | |
| # Embed the protein sequence and transpose the dimensions | |
| output_p = self.protein_embedding(input_p) | |
| output_p = output_p.transpose(1, 2) | |
| # Apply the parallel convolution and pooling layers | |
| conv_outputs = [] | |
| for conv in self.protein_convs: | |
| conv_output = conv(output_p).squeeze(-1) | |
| conv_outputs.append(conv_output) | |
| # Concatenate the convolution outputs | |
| output_p = torch.cat(conv_outputs, dim=1) | |
| else: | |
| output_p = input_p | |
| if self.protein_layers: | |
| # Apply the additional dense layers to the protein branch | |
| output_p = self.protein_branch(output_p) | |
| # Concatenate the drug and protein outputs | |
| output_t = torch.cat([output_d, output_p], dim=1) | |
| # Apply the final dense layers | |
| output_t = self.final_branch(output_t) | |
| return output_t | |