| 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: |
| |
| raise NotImplementedError |
|
|
| |
| 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) |
| |
| |
| for metric in self.metrics[stage].values(): |
| |
| metric.update(outputs.detach().float(), fitness.float()) |
| |
| if stage == "train": |
| |
| if fitness.shape[0] > 1: |
| log_dict = self.get_log_dict("train") |
| self.log_info(log_dict) |
| |
| |
| 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") |
| |
