File size: 2,352 Bytes
3b99abb | 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 | import torchmetrics
import torch
from torch.nn.functional import cross_entropy
from ..model_interface import register_model
from .base import SaprotBaseModel
@register_model
class SaprotClassificationModel(SaprotBaseModel):
def __init__(self, num_labels: int, **kwargs):
"""
Args:
num_labels: number of labels
**kwargs: other arguments for SaprotBaseModel
"""
self.num_labels = num_labels
super().__init__(task="classification", **kwargs)
def initialize_metrics(self, stage):
return {f"{stage}_acc": torchmetrics.Accuracy()}
def forward(self, inputs, coords=None):
if coords is not None:
inputs = self.add_bias_feature(inputs, coords)
# If backbone is frozen, the embedding will be the average of all residues
if self.freeze_backbone:
repr = torch.stack(self.get_hidden_states(inputs, reduction="mean"))
x = self.model.classifier.dropout(repr)
x = self.model.classifier.dense(x)
x = torch.tanh(x)
x = self.model.classifier.dropout(x)
logits = self.model.classifier.out_proj(x)
else:
logits = self.model(**inputs).logits
return logits
def loss_func(self, stage, logits, labels):
label = labels['labels']
loss = cross_entropy(logits, label)
# Update metrics
for metric in self.metrics[stage].values():
metric.update(logits.detach(), label)
if stage == "train":
log_dict = self.get_log_dict("train")
log_dict["train_loss"] = loss
self.log_info(log_dict)
# Reset train metrics
self.reset_metrics("train")
return loss
def test_epoch_end(self, outputs):
log_dict = self.get_log_dict("test")
log_dict["test_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
print(log_dict)
self.log_info(log_dict)
self.reset_metrics("test")
def validation_epoch_end(self, outputs):
log_dict = self.get_log_dict("valid")
log_dict["valid_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
self.log_info(log_dict)
self.reset_metrics("valid")
self.check_save_condition(log_dict["valid_acc"], mode="max") |