Initial release: PeptEdgeV2 (3.43M params) w/ trained weights, config, source, model card
9f16c4c verified | import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import math | |
| PHYSICOCHEMICAL_FEATURES = { | |
| 'A': [1.8, 0.0, 0.0, 89.0, 0.360, 0.76, 0.83, 0.37, 1.0, 0.0, 0.0, 0.0], | |
| 'R': [-4.5, 1.0, 1.0, 174.0, 0.293, 0.76, 0.93, 0.58, 0.0, 0.0, 1.0, 0.0], | |
| 'N': [-3.5, 0.0, 0.0, 132.0, 0.337, 0.79, 0.89, 0.54, 0.0, 0.0, 1.0, 0.0], | |
| 'D': [-3.5, -1.0, 0.0, 133.0, 0.281, 0.74, 0.72, 0.54, 0.0, 0.0, 1.0, 0.0], | |
| 'C': [2.5, 0.0, 0.0, 121.0, 0.160, 0.81, 0.87, 0.39, 1.0, 0.0, 0.0, 0.0], | |
| 'Q': [-3.5, 0.0, 0.0, 146.0, 0.316, 0.80, 0.91, 0.57, 0.0, 0.0, 1.0, 0.0], | |
| 'E': [-3.5, -1.0, 0.0, 147.0, 0.282, 0.77, 0.73, 0.60, 0.0, 0.0, 1.0, 0.0], | |
| 'G': [-0.4, 0.0, 0.0, 75.0, 0.360, 0.69, 0.75, 0.39, 1.0, 0.0, 0.0, 0.0], | |
| 'H': [-3.2, 0.1, 1.0, 155.0, 0.290, 0.80, 0.80, 0.49, 0.0, 0.0, 1.0, 0.0], | |
| 'I': [4.5, 0.0, 0.0, 131.0, 0.321, 0.79, 0.83, 0.24, 1.0, 0.0, 0.0, 0.0], | |
| 'L': [3.8, 0.0, 0.0, 131.0, 0.313, 0.78, 0.85, 0.22, 1.0, 0.0, 0.0, 0.0], | |
| 'K': [-3.9, 1.0, 0.0, 146.0, 0.329, 0.73, 0.85, 0.58, 0.0, 0.0, 1.0, 0.0], | |
| 'M': [1.9, 0.0, 0.0, 149.0, 0.302, 0.80, 0.85, 0.34, 1.0, 0.0, 0.0, 0.0], | |
| 'F': [2.8, 0.0, 0.0, 165.0, 0.287, 0.77, 0.82, 0.22, 1.0, 0.0, 0.0, 0.0], | |
| 'P': [-1.6, 0.0, 0.0, 115.0, 0.314, 0.64, 0.63, 0.42, 1.0, 0.0, 0.0, 0.0], | |
| 'S': [-0.8, 0.0, 0.0, 105.0, 0.384, 0.74, 0.74, 0.47, 0.0, 0.0, 1.0, 0.0], | |
| 'T': [-0.7, 0.0, 0.0, 119.0, 0.360, 0.76, 0.77, 0.45, 0.0, 0.0, 1.0, 0.0], | |
| 'W': [-0.9, 0.0, 0.0, 204.0, 0.277, 0.75, 0.82, 0.27, 1.0, 0.0, 0.0, 0.0], | |
| 'Y': [-1.3, 0.0, 0.0, 181.0, 0.292, 0.77, 0.79, 0.34, 0.0, 0.0, 1.0, 0.0], | |
| 'V': [4.2, 0.0, 0.0, 117.0, 0.329, 0.76, 0.85, 0.24, 1.0, 0.0, 0.0, 0.0], | |
| } | |
| AA_ORDER = 'ARNDCQEGHILKMFPSTWYV' | |
| def get_physicochem_matrix(): | |
| mat = [] | |
| for aa in AA_ORDER: | |
| mat.append(PHYSICOCHEMICAL_FEATURES[aa]) | |
| return torch.tensor(mat, dtype=torch.float32) | |
| class PhysicochemicalEmbedding(nn.Module): | |
| def __init__(self, in_dim=12, out_dim=32): | |
| super().__init__() | |
| self.proj = nn.Linear(in_dim, out_dim) | |
| self.norm = nn.LayerNorm(out_dim) | |
| def forward(self, x): | |
| return self.norm(F.gelu(self.proj(x))) | |
| class ConvBranch(nn.Module): | |
| def __init__(self, in_dim, out_dim, kernel_size, dilation=1, dropout=0.1): | |
| super().__init__() | |
| pad = dilation * (kernel_size - 1) // 2 | |
| self.conv = nn.Conv1d(in_dim, out_dim, kernel_size, padding=pad, dilation=dilation, bias=False) | |
| self.bn = nn.BatchNorm1d(out_dim) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x): | |
| x = x.transpose(1, 2) | |
| x = self.conv(x) | |
| x = self.bn(x) | |
| x = F.gelu(x) | |
| x = self.dropout(x) | |
| return x.transpose(1, 2) | |
| class SqueezeExcitation(nn.Module): | |
| def __init__(self, channels, reduction=16): | |
| super().__init__() | |
| self.fc1 = nn.Linear(channels, channels // reduction, bias=False) | |
| self.fc2 = nn.Linear(channels // reduction, channels, bias=False) | |
| def forward(self, x): | |
| b, l, c = x.shape | |
| y = x.mean(dim=1) | |
| y = F.gelu(self.fc1(y)) | |
| y = torch.sigmoid(self.fc2(y)).unsqueeze(1) | |
| return x * y | |
| class LightweightAttention(nn.Module): | |
| def __init__(self, d_model, n_heads, dropout=0.1, max_len=512): | |
| super().__init__() | |
| assert d_model % n_heads == 0 | |
| self.d_model = d_model | |
| self.n_heads = n_heads | |
| self.head_dim = d_model // n_heads | |
| self.qkv = nn.Linear(d_model, d_model * 3, bias=False) | |
| self.out = nn.Linear(d_model, d_model) | |
| self.dropout = nn.Dropout(dropout) | |
| self.pos_encoding = nn.Parameter(torch.randn(1, max_len, d_model) * 0.02) | |
| def forward(self, x, mask=None): | |
| b, l, _ = x.shape | |
| x = x + self.pos_encoding[:, :l, :] | |
| qkv = self.qkv(x).reshape(b, l, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4) | |
| q, k, v = qkv[0], qkv[1], qkv[2] | |
| scale = self.head_dim ** -0.5 | |
| attn = torch.matmul(q, k.transpose(-2, -1)) * scale | |
| if mask is not None: | |
| attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(1) == 0, float('-inf')) | |
| attn = F.softmax(attn, dim=-1) | |
| attn = self.dropout(attn) | |
| out = torch.matmul(attn, v).transpose(1, 2).contiguous().reshape(b, l, self.d_model) | |
| return self.out(out) | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, d_model, n_heads, ff_dim, dropout=0.1): | |
| super().__init__() | |
| self.norm1 = nn.LayerNorm(d_model) | |
| self.attn = LightweightAttention(d_model, n_heads, dropout) | |
| self.norm2 = nn.LayerNorm(d_model) | |
| self.ff = nn.Sequential( | |
| nn.Linear(d_model, ff_dim), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(ff_dim, d_model), | |
| nn.Dropout(dropout), | |
| ) | |
| def forward(self, x): | |
| x = x + self.attn(self.norm1(x)) | |
| x = x + self.ff(self.norm2(x)) | |
| return x | |
| class AttentionPooling(nn.Module): | |
| def __init__(self, d_model): | |
| super().__init__() | |
| self.query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02) | |
| self.attn = nn.Linear(d_model, 1) | |
| def forward(self, x): | |
| scores = self.attn(torch.tanh(x)).transpose(1, 2) | |
| weights = F.softmax(scores, dim=-1) | |
| return torch.matmul(weights, x).squeeze(1) | |
| class PeptEdge(nn.Module): | |
| def __init__(self, vocab_size=21, max_len=200, d_model=128, n_heads=4, | |
| num_layers=3, ff_dim=256, num_classes=2, dropout=0.15): | |
| super().__init__() | |
| self.max_len = max_len | |
| self.vocab_size = vocab_size | |
| phys_mat = get_physicochem_matrix() | |
| self.register_buffer('phys_mat', phys_mat) | |
| self.token_embed = nn.Embedding(vocab_size, d_model - 32, padding_idx=0) | |
| self.phys_embed = PhysicochemicalEmbedding(12, 32) | |
| conv_dims = [48, 48, 32, 32, 32] | |
| conv_ks = [3, 5, 7, 15, 31] | |
| conv_dils = [1, 1, 2, 1, 1] | |
| self.conv_branches = nn.ModuleList([ | |
| ConvBranch(d_model, conv_dims[i], conv_ks[i], conv_dils[i], dropout) | |
| for i in range(len(conv_dims)) | |
| ]) | |
| total_conv_out = sum(conv_dims) | |
| self.conv_proj = nn.Sequential( | |
| nn.Linear(total_conv_out, d_model), | |
| nn.LayerNorm(d_model), | |
| ) | |
| self.se = SqueezeExcitation(d_model) | |
| self.transformer_blocks = nn.ModuleList([ | |
| TransformerBlock(d_model, n_heads, ff_dim, dropout) | |
| for _ in range(num_layers) | |
| ]) | |
| self.attn_pool = AttentionPooling(d_model) | |
| self.classifier = nn.Sequential( | |
| nn.LayerNorm(d_model * 3), | |
| nn.Linear(d_model * 3, d_model), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(d_model, d_model // 2), | |
| nn.GELU(), | |
| nn.Dropout(dropout * 0.5), | |
| nn.Linear(d_model // 2, num_classes), | |
| ) | |
| self._init_weights() | |
| def _init_weights(self): | |
| for p in self.parameters(): | |
| if p.dim() > 1: | |
| nn.init.kaiming_normal_(p, mode='fan_out', nonlinearity='relu') | |
| def forward(self, x, return_embeddings=False): | |
| b, l = x.shape | |
| mask = (x != 0).float() | |
| aa_indices = torch.clamp(x, 0, 19).long() | |
| phys_feats = F.embedding(aa_indices, self.phys_mat) | |
| phys_emb = self.phys_embed(phys_feats) | |
| tok_emb = self.token_embed(x) | |
| h = torch.cat([tok_emb, phys_emb], dim=-1) | |
| conv_out = [] | |
| for branch in self.conv_branches: | |
| conv_out.append(branch(h)) | |
| max_len = max(c.shape[1] for c in conv_out) | |
| conv_out_padded = [] | |
| for c in conv_out: | |
| if c.shape[1] < max_len: | |
| pad = max_len - c.shape[1] | |
| c = F.pad(c, (0, 0, 0, pad)) | |
| conv_out_padded.append(c) | |
| h_conv = torch.cat(conv_out_padded, dim=-1) | |
| h_conv = self.conv_proj(h_conv) | |
| h_conv = self.se(h_conv) | |
| for block in self.transformer_blocks: | |
| h_conv = block(h_conv) | |
| mean_pool = h_conv.mean(dim=1) | |
| max_pool = h_conv.max(dim=1)[0] | |
| attn_pool = self.attn_pool(h_conv) | |
| h_pool = torch.cat([mean_pool, max_pool, attn_pool], dim=-1) | |
| if return_embeddings: | |
| return h_pool | |
| logits = self.classifier(h_pool) | |
| return logits | |
| class PeptEdgeMultilabel(nn.Module): | |
| def __init__(self, vocab_size=27, max_len=200, d_model=128, n_heads=4, | |
| num_layers=3, ff_dim=256, num_classes=5, dropout=0.15): | |
| super().__init__() | |
| self.backbone = PeptEdge( | |
| vocab_size=vocab_size, max_len=max_len, | |
| d_model=d_model, n_heads=n_heads, | |
| num_layers=num_layers, ff_dim=ff_dim, | |
| num_classes=num_classes, dropout=dropout | |
| ) | |
| in_features = d_model * 3 | |
| self.backbone.classifier = nn.Sequential( | |
| nn.LayerNorm(in_features), | |
| nn.Linear(in_features, d_model), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(d_model, d_model // 2), | |
| nn.GELU(), | |
| nn.Dropout(dropout * 0.5), | |
| nn.Linear(d_model // 2, num_classes), | |
| ) | |
| def forward(self, x, return_embeddings=False): | |
| return self.backbone(x, return_embeddings) | |
| def count_parameters(model): | |
| return sum(p.numel() for p in model.parameters() if p.requires_grad) | |
| if __name__ == '__main__': | |
| model = PeptEdge(vocab_size=21, max_len=200, num_classes=2) | |
| total = count_parameters(model) | |
| print(f'PeptEdge params: {total:,}') | |
| x = torch.randint(0, 20, (4, 100)) | |
| out = model(x) | |
| print(f'Input: {x.shape}, Output: {out.shape}') | |