| from .paths import BACKBONE |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import math |
| from transformers import AutoTokenizer |
| from torch.optim import AdamW |
| import numpy as np |
| from sklearn.metrics import f1_score |
| from typing import Tuple |
| import lightning.pytorch as pl |
| import gc |
| from torch.optim.lr_scheduler import _LRScheduler |
| from sklearn.metrics import f1_score, roc_auc_score, accuracy_score, matthews_corrcoef, precision_score, recall_score |
|
|
|
|
| class GELU(nn.Module): |
| def forward(self, x): |
| return 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0))) |
|
|
|
|
| class FocalLoss(nn.Module): |
| """Focal loss. |
| |
| Sources: |
| https://github.com/fcakyon/balanced-loss/blob/main/balanced_loss/losses.py |
| https://github.com/vandit15/Class-balanced-loss-pytorch/blob/master/class_balanced_loss.py |
| """ |
| def __init__(self, gamma=2): |
| super(FocalLoss, self).__init__() |
| self.gamma = gamma |
|
|
| def forward(self, inputs, targets, sample_weights): |
| BCLoss = F.binary_cross_entropy_with_logits(input = inputs, target = targets,reduction = "none") |
|
|
| if self.gamma == 0.0: |
| modulator = 1.0 |
| else: |
| modulator = torch.exp(-self.gamma * targets * inputs - self.gamma * torch.log(1 + |
| torch.exp(-1.0 * inputs))) |
|
|
| loss = modulator * BCLoss |
|
|
| weighted_loss = sample_weights * loss |
| focal_loss = torch.sum(weighted_loss) |
|
|
| focal_loss /= torch.sum(targets) |
| return focal_loss |
|
|
|
|
| class F1Loss(nn.Module): |
| def __init__(self, epsilon=1e-7): |
| super().__init__() |
| self.epsilon = epsilon |
|
|
| def forward(self, y_pred, y_true): |
| assert y_pred.ndim == 2 |
| assert y_true.ndim == 2 |
| y_true = y_true.float() |
| y_pred = torch.sigmoid(y_pred) |
|
|
| tp = (y_true * y_pred).sum(dim=0) |
| fp = ((1 - y_true) * y_pred).sum(dim=0) |
| fn = (y_true * (1 - y_pred)).sum(dim=0) |
|
|
| precision = tp / (tp + fp + self.epsilon) |
| recall = tp / (tp + fn + self.epsilon) |
|
|
| f1 = 2 * (precision * recall) / (precision + recall + self.epsilon) |
| f1 = f1.clamp(min=self.epsilon, max=1-self.epsilon) |
| return 1 - f1.mean() |
|
|
|
|
| def rotate_half(x): |
| x1, x2 = x.chunk(2, dim=-1) |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| def apply_rotary_pos_emb(x, cos, sin): |
| |
| cos = cos.unsqueeze(2) |
| sin = sin.unsqueeze(2) |
| return (x * cos) + (rotate_half(x) * sin) |
|
|
|
|
| class RotaryEmbedding(torch.nn.Module): |
| """Rotary position embeddings (Su et al.), inspired by GPT-NeoX. |
| |
| https://arxiv.org/abs/2104.09864 |
| https://github.com/ZhuiyiTechnology/roformer |
| https://github.com/EleutherAI/gpt-neox |
| """ |
|
|
| def __init__(self, dim: int, *_, **__): |
| super().__init__() |
| inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq) |
|
|
| self._seq_len_cached = None |
| self._cos_cached = None |
| self._sin_cached = None |
|
|
| def _update_cos_sin_tables(self, x, seq_dimension=1): |
| seq_len = x.shape[seq_dimension] |
|
|
| if seq_len != self._seq_len_cached or self._cos_cached.device != x.device: |
| self._seq_len_cached = seq_len |
| t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq) |
| freqs = torch.einsum("i,j->ij", t, self.inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1).to(x.device) |
|
|
| self._cos_cached = emb.cos()[None, :, :] |
| self._sin_cached = emb.sin()[None, :, :] |
|
|
| return self._cos_cached, self._sin_cached |
|
|
| def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k) |
|
|
| return ( |
| apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), |
| apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached), |
| ) |
|
|
|
|
| def macro_f1(y_true, y_pred, thresholds): |
| y_pred_binary = (y_pred >= thresholds).float() |
| tp = (y_true * y_pred_binary).sum(dim=0) |
| fp = ((1 - y_true) * y_pred_binary).sum(dim=0) |
| fn = (y_true * (1 - y_pred_binary)).sum(dim=0) |
| precision = tp / (tp + fp + 1e-7) |
| recall = tp / (tp + fn + 1e-7) |
| f1 = 2 * precision * recall / (precision + recall + 1e-7) |
| macro_f1 = f1.mean() |
| return macro_f1 |
|
|
|
|
| def safeguard_softmax(logits, dim=-1): |
| max_logits, _ = logits.max(dim=dim, keepdim=True) |
| exp_logits = torch.exp(logits - max_logits) |
| exp_sum = exp_logits.sum(dim=dim, keepdim=True) |
| probs = exp_logits / (exp_sum + 1e-7) |
| return probs |
|
|
|
|
| class PositionalAttentionHead(nn.Module): |
| def __init__(self, hidden_dim, n_heads): |
| super(PositionalAttentionHead, self).__init__() |
| self.n_heads = n_heads |
| self.hidden_dim = hidden_dim |
| self.head_dim = hidden_dim // n_heads |
| self.preattn_ln = nn.LayerNorm(self.head_dim) |
| self.Q = nn.Linear(self.head_dim, self.head_dim, bias=False) |
| self.K = nn.Linear(self.head_dim, self.head_dim, bias=False) |
| self.V = nn.Linear(self.head_dim, self.head_dim, bias=False) |
| self.rot_emb = RotaryEmbedding(self.head_dim) |
|
|
| def forward(self, x, attention_mask): |
| batch_size, seq_len, _ = x.size() |
| x = x.view(batch_size, seq_len, self.n_heads, self.head_dim) |
| x = self.preattn_ln(x) |
|
|
| q = self.Q(x) |
| k = self.K(x) |
| v = self.V(x) |
|
|
| q, k = self.rot_emb(q, k) |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
| attn_scores = torch.einsum('bqhd,bkhd->bhqk', q, k) / math.sqrt(self.head_dim) |
| attn_scores = attn_scores.masked_fill(torch.logical_not(attention_mask.unsqueeze(1).unsqueeze(1)), float("-inf")) |
|
|
| attn_probs = safeguard_softmax(attn_scores, dim=-1) |
|
|
| x = torch.einsum('bhqk,bkhd->bqhd', attn_probs, v) |
| x = x.reshape(batch_size, seq_len, self.hidden_dim) |
| gc.collect() |
| torch.cuda.empty_cache() |
| return x, attn_probs |
|
|
|
|
| class CosineAnnealingWithWarmup(_LRScheduler): |
| def __init__(self, optimizer, warmup_steps, total_steps, eta_ratio=0.1, last_epoch=-1): |
| self.warmup_steps = warmup_steps |
| self.total_steps = total_steps |
| self.eta_ratio = eta_ratio |
| super(CosineAnnealingWithWarmup, self).__init__(optimizer, last_epoch) |
|
|
| def get_lr(self): |
| if self.last_epoch < self.warmup_steps: |
| return [base_lr * self.last_epoch / self.warmup_steps for base_lr in self.base_lrs] |
|
|
| progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps) |
| cosine_decay = 0.5 * (1 + np.cos(np.pi * progress)) |
| decayed_lr = (1 - self.eta_ratio) * cosine_decay + self.eta_ratio |
|
|
| return [decayed_lr * base_lr for base_lr in self.base_lrs] |
|
|
|
|
| class AdaptiveMarginTripletLoss(nn.Module): |
| def __init__(self, initial_margin=0.05, scale_factor=0.1, min_margin=0.05, max_margin=2.0): |
| super(AdaptiveMarginTripletLoss, self).__init__() |
| self.initial_margin = initial_margin |
| self.scale_factor = scale_factor |
| self.min_margin = min_margin |
| self.max_margin = max_margin |
|
|
| def forward(self, anchor, positive, negative): |
| pos_dist = F.pairwise_distance(anchor, positive, p=2) |
| neg_dist = F.pairwise_distance(anchor, negative, p=2) |
|
|
| adaptive_margins = self.initial_margin + self.scale_factor * (neg_dist - pos_dist) |
|
|
| adaptive_margins = torch.clamp(adaptive_margins, min=self.min_margin, max=self.max_margin) |
|
|
| losses = F.relu(pos_dist - neg_dist + adaptive_margins) |
|
|
| return losses.mean() |
|
|
|
|
| class RobertaLMHead(nn.Module): |
| def __init__(self, embed_dim, output_dim, weight): |
| super().__init__() |
| self.dense = nn.Linear(embed_dim, embed_dim) |
| self.layer_norm = nn.LayerNorm(embed_dim) |
| self.weight = weight |
| self.gelu = GELU() |
| self.bias = nn.Parameter(torch.zeros(output_dim)) |
| def forward(self, features): |
| x = self.dense(features) |
| x = self.gelu(x) |
| x = self.layer_norm(x) |
| x = F.linear(x, self.weight) + self.bias |
| return x |
|
|
|
|
| def pad_to_max_length(tensor, max_length, pad_value=1): |
| return F.pad(tensor, (0, max_length - tensor.shape[1]), value=pad_value) |
|
|
|
|
| def compute_ece(probs, labels, n_bins=10): |
| """Per-class calibration error for probabilities and labels of shape (B, C).""" |
| bin_boundaries = torch.linspace(0, 1, n_bins + 1) |
| bin_lowers = bin_boundaries[:-1] |
| bin_uppers = bin_boundaries[1:] |
|
|
| n_classes = probs.size(1) |
| ece_per_class = torch.zeros(n_classes, device=probs.device) |
|
|
| for class_idx in range(n_classes): |
| class_probs = probs[:, class_idx] |
| class_labels = labels[:, class_idx] |
|
|
| ece = torch.zeros(1, device=probs.device) |
| for bin_lower, bin_upper in zip(bin_lowers, bin_uppers): |
| in_bin = (class_probs.gt(bin_lower.item()) * class_probs.le(bin_upper.item())).float() |
| prop_in_bin = in_bin.mean() |
|
|
| if prop_in_bin > 0: |
| accuracy_in_bin = (class_labels * in_bin).sum() / in_bin.sum() |
| avg_confidence_in_bin = (class_probs * in_bin).sum() / in_bin.sum() |
| ece += torch.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin |
|
|
| ece_per_class[class_idx] = ece.item() |
|
|
| return ece_per_class |
|
|
|
|
| class MultitaskProteinModel(pl.LightningModule): |
| def __init__(self, hparams, esm_model, num_layers_to_finetune, input_dim=1280, output_dim=12, class_weights=None, hidden_dim=512, n_heads=10): |
| super(MultitaskProteinModel, self).__init__() |
| self.save_hyperparameters(hparams) |
| self.esm_model = esm_model |
| for param in self.esm_model.parameters(): |
| param.requires_grad = False |
| for i in range(num_layers_to_finetune): |
| for param in self.esm_model.encoder.layer[-i-1].parameters(): |
| param.requires_grad = True |
|
|
| self.lm_head = RobertaLMHead(embed_dim = 1280, output_dim=33, weight=self.esm_model.embeddings.word_embeddings.weight) |
| self.attn_head = PositionalAttentionHead(input_dim, n_heads) |
| self.attn_ln = nn.LayerNorm(input_dim) |
| self.attn_skip = nn.Linear(input_dim, input_dim) |
| self.linear_layers = nn.ModuleList() |
| for _ in range(self.hparams.num_linear_layers): |
| self.linear_layers.append(nn.Linear(input_dim, input_dim)) |
| self.reduction_layers = nn.Sequential( |
| nn.Linear(input_dim, hidden_dim), |
| GELU(), |
| nn.Linear(hidden_dim, output_dim) |
| ) |
| self.clf_ln = nn.LayerNorm(input_dim) |
| self.class_weights = class_weights |
| self.classification_thresholds = nn.Parameter(torch.tensor([0.5]*output_dim)) |
| self.historical_memory = 0.9 |
| self.training_step_outputs = [] |
| self.training_step_true_labels = [] |
| self.testing_step_predictions = [] |
| self.testing_step_labels = [] |
| self.test_dataloader = None |
| self.tokenizer = AutoTokenizer.from_pretrained(str(BACKBONE), local_files_only=True) |
| self.corruption_epoch_threshold = 5 |
|
|
| def setup(self, stage=None): |
| self.__build_loss() |
|
|
| def __build_loss(self): |
| self.contrastive_loss = AdaptiveMarginTripletLoss() |
| focal_gamma = torch.tensor(self.hparams.focal_gamma).to(self.device) |
| self.class_weights = self.class_weights.to(self.device) |
| self.focal_loss = FocalLoss(gamma=focal_gamma) |
| self.f1_loss = F1Loss() |
| self.reconstruction_loss = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id) |
|
|
| def forward(self, input_ids, attention_masks=None): |
| outputs = self.esm_model(input_ids=input_ids, attention_mask=attention_masks, output_hidden_states=True) |
| embeddings = outputs.last_hidden_state |
| attention_masks = attention_masks |
|
|
| x_pool, x_attns = self.attn_head(embeddings, attention_masks) |
| x_pool = self.attn_ln(x_pool + self.attn_skip(x_pool)) |
|
|
| for linear_layer in self.linear_layers: |
| residue = x_pool |
| x_pool = linear_layer(x_pool) |
| x_pool = F.silu(x_pool) |
| x_pool = x_pool + residue |
|
|
| x_weighted = torch.einsum('bhlk,bld->bhld', x_attns, x_pool) |
| x_combined = x_weighted.mean(dim=1) |
| x_combined = self.clf_ln(x_combined) |
| mlm_logits = self.lm_head(x_combined) |
| attention_masks = attention_masks.unsqueeze(-1).float() |
| attention_sum = attention_masks.sum(dim=1, keepdim=True) |
| x_combined_masked = (x_combined * attention_masks).sum(dim=1) / attention_sum.squeeze(1) |
|
|
| x_pred = self.reduction_layers(x_combined_masked) |
| gc.collect() |
| torch.cuda.empty_cache() |
| return x_pred, x_attns, x_combined_masked, mlm_logits |
| def forward_from_embeds(self, inputs_embeds, attention_masks=None): |
| esm = self.esm_model |
| |
| if getattr(esm.embeddings, "token_dropout", False): |
| esm.embeddings.token_dropout = False |
|
|
| outs = esm( |
| inputs_embeds=inputs_embeds, |
| attention_mask=attention_masks.to(torch.long), |
| output_hidden_states=True, |
| ) |
| embeddings = outs.last_hidden_state |
|
|
| x_pool, x_attns = self.attn_head(embeddings, attention_masks) |
| x_pool = self.attn_ln(x_pool + self.attn_skip(x_pool)) |
| for linear_layer in self.linear_layers: |
| residue = x_pool |
| x_pool = F.silu(linear_layer(x_pool)) |
| x_pool = x_pool + residue |
|
|
| x_weighted = torch.einsum('bhlk,bld->bhld', x_attns, x_pool).mean(dim=1) |
| x_weighted = self.clf_ln(x_weighted) |
|
|
| mask_f = attention_masks.unsqueeze(-1).float() |
| x_combined_masked = (x_weighted * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp_min(1e-6) |
|
|
| x_pred = self.reduction_layers(x_combined_masked) |
| return x_pred, x_attns, x_combined_masked, None |
| def predict(self, input_ids, attention_masks=None): |
| x_pred, _, _, _ = self.forward(input_ids, attention_masks) |
| classification_output = torch.sigmoid(x_pred) |
| predictions = (classification_output >= self.classification_thresholds).float() |
|
|
| for i, pred in enumerate(predictions): |
| if pred.sum() == 0: |
| weighted_probs = classification_output[i] |
| max_class = torch.argmax(weighted_probs) |
| predictions[i, max_class] = 1.0 |
|
|
| confidence_scores = torch.abs(classification_output - self.classification_thresholds) / self.classification_thresholds |
| max_confidence_scores = confidence_scores.max(dim=1).values |
|
|
| return classification_output, predictions, max_confidence_scores |
|
|
| def training_step(self, batch, batch_idx): |
| anchor_input_ids, anchor_attention_masks = batch['anchor_input_ids'], batch['anchor_attention_mask'] |
| positive_input_ids, positive_attention_masks = batch['positive_input_ids'], batch['positive_attention_mask'] |
| negative_input_ids, negative_attention_masks = batch['negative_input_ids'], batch['negative_attention_mask'] |
| anchor_labels = batch['anchor_labels'] |
| positive_labels = batch['positive_labels'] |
| negative_labels = batch['negative_labels'] |
|
|
| anchor_x_pred, _, anchor_x_combined, mlm_logits = self.forward(anchor_input_ids, anchor_attention_masks) |
| positive_x_pred, _, positive_x_combined, _ = self.forward(positive_input_ids, positive_attention_masks) |
| negative_x_pred, _, negative_x_combined, _ = self.forward(negative_input_ids, negative_attention_masks) |
|
|
| reconstruction_loss = self.reconstruction_loss(mlm_logits.view(-1, mlm_logits.size(-1)), anchor_input_ids.view(-1)) |
|
|
| anchor_focal_loss = self.focal_loss(anchor_x_pred, anchor_labels, self.class_weights) |
| positive_focal_loss = self.focal_loss(positive_x_pred, positive_labels, self.class_weights) |
|
|
| semi_hard_negative_ce_loss = F.binary_cross_entropy_with_logits(input=negative_x_pred, target=negative_labels) |
|
|
| focal_loss = (anchor_focal_loss + positive_focal_loss) / 2 |
|
|
| anchor_f1_loss = self.f1_loss(anchor_x_pred, anchor_labels) |
| positive_f1_loss = self.f1_loss(positive_x_pred, positive_labels) |
| semi_hard_f1_loss = self.f1_loss(negative_x_pred, negative_labels) |
|
|
| semi_hard_triplet_loss = self.contrastive_loss(anchor_x_combined, positive_x_combined, negative_x_combined) |
|
|
| if self.current_epoch > self.corruption_epoch_threshold: |
| hard_negative_input_ids, hard_negative_attention_masks = batch['hard_negative_input_ids'], batch['hard_negative_attention_mask'] |
| hard_negative_labels = batch['hard_negative_labels'] |
| if self.current_epoch < 15: |
| epoch_ratio = min(1.0, (self.current_epoch - self.corruption_epoch_threshold) / (self.hparams.epochs - self.corruption_epoch_threshold)) |
| hard_weight = epoch_ratio |
| semi_hard_weight = 1 - hard_weight |
| else: |
| semi_hard_weight = 0.0 |
| hard_weight = 1.0 |
| hard_negative_x_pred, _, hard_negative_x_combined, _ = self.forward(hard_negative_input_ids, hard_negative_attention_masks) |
| hard_negative_ce_loss = F.binary_cross_entropy_with_logits(input=hard_negative_x_pred, target=hard_negative_labels) |
| hard_triplet_loss = self.contrastive_loss(anchor_x_combined, positive_x_combined, hard_negative_x_combined) |
| hard_f1_loss = self.f1_loss(hard_negative_x_pred, hard_negative_labels) |
| negative_ce_loss = semi_hard_weight * semi_hard_negative_ce_loss + hard_weight * hard_negative_ce_loss |
| triplet_loss = semi_hard_weight * semi_hard_triplet_loss + hard_weight * hard_triplet_loss |
| f1_loss = anchor_f1_loss + positive_f1_loss + semi_hard_weight * semi_hard_triplet_loss + hard_weight * hard_f1_loss |
| else: |
| negative_ce_loss = semi_hard_negative_ce_loss |
| triplet_loss = semi_hard_triplet_loss |
| f1_loss = anchor_f1_loss + positive_f1_loss + semi_hard_f1_loss |
|
|
| loss = focal_loss * 2 + negative_ce_loss + f1_loss + triplet_loss * 2 + reconstruction_loss |
|
|
| self.training_step_outputs.append(anchor_x_pred) |
| self.training_step_outputs.append(positive_x_pred) |
| self.training_step_outputs.append(negative_x_pred) |
| self.training_step_true_labels.append(anchor_labels) |
| self.training_step_true_labels.append(positive_labels) |
| self.training_step_true_labels.append(negative_labels) |
| macro_f1_val = macro_f1(anchor_labels, torch.sigmoid(anchor_x_pred), self.classification_thresholds) |
| self.log('train_macro_f1', macro_f1_val, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) |
| gc.collect() |
| torch.cuda.empty_cache() |
| gc.collect() |
| torch.cuda.empty_cache() |
| self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) |
| return loss |
|
|
| def on_train_epoch_end(self): |
| epoch_outputs = torch.cat(self.training_step_outputs, dim=0) |
| epoch_true_labels = torch.cat(self.training_step_true_labels, dim=0) |
| epoch_outputs = torch.sigmoid(epoch_outputs) |
| self.update_class_thresholds(epoch_outputs, epoch_true_labels) |
| threshold_dict = {f"threshold_{i}": self.classification_thresholds[i].item() for i in range(15)} |
| self.log_dict(threshold_dict, on_epoch=True) |
| self.training_step_outputs.clear() |
| self.training_step_true_labels.clear() |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
| def validation_step(self, batch, batch_idx): |
| anchor_input_ids, anchor_attention_masks = batch['anchor_input_ids'], batch['anchor_attention_mask'] |
| positive_input_ids, positive_attention_masks = batch['positive_input_ids'], batch['positive_attention_mask'] |
| negative_input_ids, negative_attention_masks = batch['negative_input_ids'], batch['negative_attention_mask'] |
| anchor_labels = batch['anchor_labels'] |
| positive_labels = batch['positive_labels'] |
| negative_labels = batch['negative_labels'] |
|
|
| anchor_x_pred, _, anchor_x_combined, mlm_logits = self.forward(anchor_input_ids, anchor_attention_masks) |
| positive_x_pred, _, positive_x_combined, _ = self.forward(positive_input_ids, positive_attention_masks) |
| negative_x_pred, _, negative_x_combined, _ = self.forward(negative_input_ids, negative_attention_masks) |
|
|
| reconstruction_loss = self.reconstruction_loss(mlm_logits.view(-1, mlm_logits.size(-1)), anchor_input_ids.view(-1)) |
| anchor_focal_loss = self.focal_loss(anchor_x_pred, anchor_labels , self.class_weights) |
| positive_focal_loss = self.focal_loss(positive_x_pred, positive_labels, self.class_weights) |
|
|
| negative_ce_loss = F.binary_cross_entropy_with_logits(input=negative_x_pred, target=negative_labels) |
|
|
| focal_loss = (anchor_focal_loss + positive_focal_loss) / 2 |
|
|
| f1_loss = self.f1_loss(anchor_x_pred, anchor_labels) |
|
|
| triplet_loss = self.contrastive_loss(anchor_x_combined, positive_x_combined, negative_x_combined) |
|
|
| loss = focal_loss * 2 + negative_ce_loss + f1_loss + triplet_loss * 2 + reconstruction_loss |
|
|
| macro_f1_val = macro_f1(anchor_labels, torch.sigmoid(anchor_x_pred), self.classification_thresholds) |
| self.log('val_macro_f1', macro_f1_val, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) |
| gc.collect() |
| torch.cuda.empty_cache() |
| self.log('val_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) |
|
|
|
|
| def test_step(self, batch, batch_idx): |
| input_ids, labels, attention_masks = batch['input_ids'], batch['labels'], batch['attention_masks'] |
| x_pred, _, _ = self.forward_test(input_ids, attention_masks) |
| predictions = torch.sigmoid(x_pred) |
| self.testing_step_predictions.append(predictions) |
| self.testing_step_labels.append(labels) |
|
|
| def on_test_epoch_end(self,): |
| predictions = torch.cat(self.testing_step_predictions, dim=0) |
| labels = torch.cat(self.testing_step_labels, dim=0) |
|
|
| predictions_binary = (predictions >= self.classification_thresholds).float().cpu().numpy() |
| labels_binary = labels.cpu().numpy() |
|
|
| f1 = f1_score(labels_binary, predictions_binary, average='macro') |
| roc_auc = roc_auc_score(labels_binary, predictions.cpu().to(torch.float32).numpy(), average='micro') |
| accuracy = accuracy_score(labels_binary.flatten(), predictions_binary.flatten()) |
| mcc = matthews_corrcoef(labels_binary.flatten(), predictions_binary.flatten()) |
| precision = precision_score(labels_binary, predictions_binary, average='macro') |
| recall = recall_score(labels_binary, predictions_binary, average='macro') |
|
|
| results = { |
| 'test_F1_macro': f1, |
| 'test_AUCROC': roc_auc, |
| 'test_Accuracy': accuracy, |
| 'test_MCC': mcc, |
| 'test_Precision': precision, |
| 'test_Recall': recall, |
| } |
| self.log_dict(results) |
| self.testing_step_predictions.clear() |
| self.testing_step_labels.clear() |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
|
|
| def configure_optimizers(self): |
| optimizer = AdamW(self.parameters(), lr=self.hparams.lr, betas=(0.9, 0.95), weight_decay=1e-5) |
| schedulers = CosineAnnealingWithWarmup(optimizer, warmup_steps=100, eta_ratio=0.1, total_steps=9000) |
| lr_schedulers = { |
| "scheduler": schedulers, |
| "name": 'learning_rate_logs', |
| "interval": 'step', |
| 'frequency': 1 |
| } |
| return [optimizer], [lr_schedulers] |
|
|
|
|
| def update_class_thresholds(self, inputs, targets): |
| with torch.no_grad(): |
| min_threshold_value = 0.001 |
| thresholds = torch.linspace(0.0, 1.0, steps=1000, device=inputs.device) |
| num_classes = targets.size(1) |
|
|
| best_f1_scores = torch.zeros(num_classes, device=inputs.device) |
| best_thresholds = torch.full((num_classes,), min_threshold_value, device=inputs.device) |
|
|
| for threshold in thresholds: |
| binary_predictions = (inputs >= threshold).float() |
|
|
| tp = (binary_predictions * targets).sum(dim=0) |
| fp = (binary_predictions * (1 - targets)).sum(dim=0) |
| fn = ((1 - binary_predictions) * targets).sum(dim=0) |
|
|
| precision = tp / (tp + fp + 1e-7) |
| recall = tp / (tp + fn + 1e-7) |
| f1_scores = 2 * precision * recall / (precision + recall + 1e-7) |
|
|
| is_better = f1_scores > best_f1_scores |
| best_f1_scores = torch.where(is_better, f1_scores, best_f1_scores) |
| best_thresholds = torch.where(is_better, threshold.expand_as(best_f1_scores), best_thresholds) |
|
|
| updated_thresholds = self.historical_memory * self.classification_thresholds.data + (1 - self.historical_memory) * best_thresholds |
|
|
| self.classification_thresholds = nn.Parameter(torch.clamp(updated_thresholds, min=min_threshold_value)) |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
|
|
| class ResidueEncoder(MultitaskProteinModel): |
| """Return per-residue clf_ln outputs for diffusion training.""" |
| def forward(self, input_ids, attention_masks=None): |
| captured = [] |
| hook = self.clf_ln.register_forward_hook(lambda module, inputs, output: captured.append(output)) |
| try: |
| logits, attention, pooled, mlm_logits = super().forward(input_ids, attention_masks) |
| finally: |
| hook.remove() |
| return logits, attention, captured[0], mlm_logits |
|
|
