Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 3,595 Bytes
22761bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import torch.nn as nn
import torch.nn.functional as F
class ElectraDTA(nn.Module):
def __init__(self, smilen, seq_len, hidden_dim):
super().__init__()
self.drug_input = nn.Linear(smilen, hidden_dim)
self.prot_input = nn.Linear(seq_len, hidden_dim)
self.num_filters = hidden_dim
self.filter_length = 3
self.n_layers = 4
self.seblock = True
self.encode_smiles = ConvBlock(hidden_dim, self.seblock, self.num_filters, self.filter_length)
self.encode_prot = ConvBlock(hidden_dim, self.seblock, self.num_filters, self.filter_length)
self.global_pool = nn.AdaptiveMaxPool1d(1)
self.concat = nn.Linear(hidden_dim * 4, hidden_dim * 4)
self.predictions = FCNet(hidden_dim * 4) # you need to define this function
self.interaction_model = nn.Sequential(self.drug_input, self.prot_input, self.encode_smiles, self.encode_prot,
self.global_pool, self.concat, self.predictions)
def forward(self, x):
x = self.interaction_model(x)
return x
class SEBlock(nn.Module):
def __init__(self, channels, r=8):
super().__init__()
self.squeeze = nn.AdaptiveAvgPool1d(1)
self.excitation = nn.Sequential(
nn.Linear(channels, channels // r),
nn.ReLU(),
nn.Linear(channels // r, channels),
nn.Sigmoid()
)
def forward(self, x):
out = self.squeeze(x)
out = self.excitation(out)
return x * out
class ConvBlock(nn.Module):
def __init__(self, input_channels, seblock, num_filters, filter_length):
super().__init__()
self.conv1 = nn.Conv1d(input_channels, num_filters, filter_length, padding='valid', stride=1)
self.conv2 = nn.Conv1d(num_filters, num_filters * 2, filter_length, padding='valid', stride=1)
self.seblock = seblock
if seblock:
self.se1 = SEBlock(num_filters)
self.se2 = SEBlock(num_filters * 2)
def forward(self, x):
x = F.relu(self.conv1(x))
if self.seblock:
x = self.se1(x)
x = F.relu(self.conv2(x))
if self.seblock:
x = self.se2(x)
return x
class Highway(nn.Module):
def __init__(self, dim, n_layers, activation=nn.Tanh(), gate_bias=0):
super(Highway, self).__init__()
self.n_layers = n_layers
self.activation = activation
self.T_gates = nn.ModuleList([nn.Linear(dim, dim) for _ in range(n_layers)])
self.transforms = nn.ModuleList([nn.Linear(dim, dim) for _ in range(n_layers)])
self.sigmoid = nn.Sigmoid()
nn.init.constant_(self.linear.bias, gate_bias)
def forward(self, x):
for i in range(self.n_layers):
T = self.sigmoid(self.gate(x))
C = 1 - T
H = self.activation(self.linear(x))
x = T * H + C * x
return x
class FCNet(nn.Module):
def __init__(self, input_dim):
super().__init__()
self.n_layers = 4
self.highway = Highway(input_dim, self.n_layers, gate_bias=-2)
self.fc1 = nn.Linear(input_dim, 1024)
self.fc2 = nn.Linear(1024, 1024)
self.fc3 = nn.Linear(1024, 512)
# self.fc4 = nn.Linear(512, 1)
self.dropout = nn.Dropout(0.4)
def forward(self, x):
x = self.highway(x)
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = F.relu(self.fc2(x))
x = self.dropout(x)
x = F.relu(self.fc3(x))
# x = self.fc4(x)
return x
|