File size: 3,584 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import torch.distributed as dist
import torchmetrics
import torch
from ..model_interface import register_model
from .base import SaprotBaseModel
@register_model
class SaprotRegressionModel(SaprotBaseModel):
def __init__(self, test_result_path: str = None, **kwargs):
"""
Args:
test_result_path: path to save test result
**kwargs: other arguments for SaprotBaseModel
"""
self.test_result_path = test_result_path
super().__init__(task="regression", **kwargs)
def initialize_metrics(self, stage):
return {f"{stage}_loss": torchmetrics.MeanSquaredError(),
f"{stage}_spearman": torchmetrics.SpearmanCorrCoef(),
f"{stage}_R2": torchmetrics.R2Score(),
f"{stage}_pearson": torchmetrics.PearsonCorrCoef()}
def forward(self, inputs, structure_info=None):
if structure_info:
# To be implemented
raise NotImplementedError
# 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).squeeze(dim=-1)
else:
logits = self.model(**inputs).logits.squeeze(dim=-1)
return logits
def loss_func(self, stage, outputs, labels):
fitness = labels['labels'].to(outputs)
loss = torch.nn.functional.mse_loss(outputs, fitness)
# Update metrics
for metric in self.metrics[stage].values():
# Training is on half precision, but metrics expect float to compute correctly.
metric.update(outputs.detach().float(), fitness.float())
if stage == "train":
# Skip calculating metrics if the batch size is 1
if fitness.shape[0] > 1:
log_dict = self.get_log_dict("train")
self.log_info(log_dict)
# Reset train metrics
self.reset_metrics("train")
return loss
def test_epoch_end(self, outputs):
if self.test_result_path is not None:
from torchmetrics.utilities.distributed import gather_all_tensors
preds = self.test_spearman.preds
preds[-1] = preds[-1].unsqueeze(dim=0) if preds[-1].shape == () else preds[-1]
preds = torch.cat(gather_all_tensors(torch.cat(preds, dim=0)))
targets = self.test_spearman.target
targets[-1] = targets[-1].unsqueeze(dim=0) if targets[-1].shape == () else targets[-1]
targets = torch.cat(gather_all_tensors(torch.cat(targets, dim=0)))
if dist.get_rank() == 0:
with open(self.test_result_path, 'w') as w:
w.write("pred\ttarget\n")
for pred, target in zip(preds, targets):
w.write(f"{pred.item()}\t{target.item()}\n")
log_dict = self.get_log_dict("test")
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")
self.log_info(log_dict)
self.reset_metrics("valid")
self.check_save_condition(log_dict["valid_loss"], mode="min")
|