File size: 1,531 Bytes
eda7ec8 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
class CNNED_Protein(nn.Module):
def __init__(self, alphabet_size: int, target_size: int,
channel: int, depth: int, kernel_size: int, l2norm: bool = True):
super().__init__()
C_in = alphabet_size
C = channel
K = kernel_size
pad = K // 2
blocks = [
nn.Conv1d(C_in, C, K, stride=1, padding=pad, bias=False),
nn.BatchNorm1d(C),
nn.ReLU(inplace=True),
]
for _ in range(depth - 1):
blocks += [
nn.Conv1d(C, C, K, stride=1, padding=pad, bias=False),
nn.BatchNorm1d(C),
nn.ReLU(inplace=True),
nn.AvgPool1d(2),
]
self.conv = nn.Sequential(*blocks)
self.pool = nn.AdaptiveAvgPool1d(1)
self.proj = nn.Sequential(
nn.Linear(C, C),
nn.ReLU(inplace=True),
nn.Linear(C, target_size),
)
self.l2norm = l2norm
def encode(self, x: torch.Tensor):
# x: (B, A, L)
z = self.conv(x) # (B, C, L')
z = self.pool(z).squeeze(-1) # (B, C)
y = self.proj(z) # (B, D)
if self.l2norm:
y = F.normalize(y, dim=-1)
return y, z
def forward(self, a: torch.Tensor, p: torch.Tensor, n: torch.Tensor):
ay, _ = self.encode(a)
py, _ = self.encode(p)
ny, _ = self.encode(n)
return ay, py, ny
|