device-agnostic model loading + formatting

Author: jannisborn

scripts/train_tox.py

@@ -86,7 +86,6 @@ def main(
     embedding_path: str,
     finetune_path: str,
 ):
-
     logging.basicConfig(level=logging.INFO, format="%(message)s")
     logger = logging.getLogger(f"{training_name}")
     logger.setLevel(logging.INFO)
@@ -282,7 +281,7 @@ def smiles_tensor_batch_to_fp(smiles):
     if finetune_path:
         if os.path.isfile(finetune_path):
             try:
-                model.load(finetune_path)
+                model.load(finetune_path, map_location=device)
                 logger.info(f"Restored pretrained model {finetune_path}")
             except Exception:
                 raise KeyError(f"Could not restore model from {finetune_path}")
@@ -313,7 +312,6 @@ def smiles_tensor_batch_to_fp(smiles):
     )
 
     for epoch in range(params["epochs"]):
-
         performer.epoch += 1
         model.train()
         logger.info(params_filepath.split("/")[-1])
@@ -349,7 +347,6 @@ def smiles_tensor_batch_to_fp(smiles):
             predictions = []
             labels = []
             for ind, (smiles, y) in enumerate(test_loader):
-
                 smiles = torch.squeeze(smiles.to(device))
                 # Transform smiles to FP if needed
                 if params.get("model_fn", "mca") == "dense":

toxsmi/utils/performance.py

@@ -56,9 +56,8 @@ def __init__(
         train_batches: int,
         test_batches: int,
         task_names: List[str],
-        beta: float=1
+        beta: float = 1,
     ):
-
         if task == "binary_classification":
             self.report = self.performance_report_binary_classification
             self.inference_report = self.inference_report_binary_classification
@@ -70,14 +69,14 @@ def __init__(
         elif task == "regression":
             self.report = self.performance_report_regression
             self.inference_report = self.inference_report_regression
-            self.metric_initializer("rmse", 10 ** 9)
-            self.metric_initializer("mae", 10 ** 9)
+            self.metric_initializer("rmse", 10**9)
+            self.metric_initializer("mae", 10**9)
             self.metric_initializer("pearson", -1)
             self.metric_initializer("spearman", -1)
             self.task_final_report = self.final_report_regression
         else:
             raise ValueError(f"Unknown task {task}")
-        self.metric_initializer("loss", 10 ** 9)
+        self.metric_initializer("loss", 10**9)
 
         self.task = task
         self.task_names = task_names
@@ -130,7 +129,9 @@ def performance_report_binary_classification(
         negative_recall = report["0.0"]["recall"]
         positive_precision = report["1.0"]["precision"]
         positive_recall = report["1.0"]["recall"]
-        f1 = fbeta_score(labels, bin_preds, beta=self.beta,pos_label=1, average='binary')
+        f1 = fbeta_score(
+            labels, bin_preds, beta=self.beta, pos_label=1, average="binary"
+        )
 
         logger.info(
             f"\t **** TEST **** Epoch [{self.epoch + 1}/{self.epochs}], "
@@ -148,9 +149,9 @@ def performance_report_binary_classification(
             "test_auc": roc_auc,
             "best_auc": self.roc_auc,
             "test_precision_recall": precision_recall,
-            'best_precision_recall': self.precision_recall,
+            "best_precision_recall": self.precision_recall,
             "test_f1": f1,
-            'best_f1': self.f1
+            "best_f1": self.f1,
         }
         self.metrics.append(info)
         if roc_auc > self.roc_auc:
@@ -163,8 +164,8 @@ def performance_report_binary_classification(
             best = "Precision-Recall"
         if f1 > self.f1:
             self.f1 = f1
-            self.save_model(model, 'F1', 'best', value=f1)
-            best = 'F1'
+            self.save_model(model, "F1", "best", value=f1)
+            best = "F1"
         if loss_a < self.loss:
             self.loss = loss_a
             self.save_model(model, "loss", "best", value=loss_a)
@@ -236,13 +237,13 @@ def inference_report_binary_classification(
         precision, recall, _ = precision_recall_curve(labels, preds)
         precision_recall = average_precision_score(labels, preds)
         report = classification_report(labels, bin_preds, output_dict=True)
-        negative_precision = report["0"]["precision"]
-        negative_recall = report["0"]["recall"]
-        positive_precision = report["1"]["precision"]
-        positive_recall = report["1"]["recall"]
+        negative_precision = report["0.0"]["precision"]
+        negative_recall = report["0.0"]["recall"]
+        positive_precision = report["1.0"]["precision"]
+        positive_recall = report["1.0"]["recall"]
         accuracy = accuracy_score(labels, bin_preds)
         bal_accuracy = balanced_accuracy_score(labels, bin_preds)
-        f1 = fbeta_score(labels, bin_preds, beta=0.5,pos_label=1, average='binary')
+        f1 = fbeta_score(labels, bin_preds, beta=0.5, pos_label=1, average="binary")
 
         info = {
             "roc_auc": roc_auc,
@@ -281,7 +282,7 @@ def inference_report_binarized_regression(
         positive_recall = report.get("1", {"recall": 0.0})["recall"]
         accuracy = accuracy_score(bin_labels, bin_preds)
         bal_accuracy = balanced_accuracy_score(bin_labels, bin_preds)
-        f1 = fbeta_score(bin_labels, bin_preds, beta=0.5 ,pos_label=1, average='binary')
+        f1 = fbeta_score(bin_labels, bin_preds, beta=0.5, pos_label=1, average="binary")
 
         info = {
             "f1": f1,