add tensorboard option and update RNA_CM_exp

Author: wisskarrou

base/RNA_CM_exp.py

@@ -8,31 +8,89 @@
 
 from rnaglib.learning.task_models import PygModel
 from rnaglib.tasks import get_task, RNA_CM
-from rnaglib.transforms import GraphRepresentation
+from rnaglib.transforms import GraphRepresentation, RNAFMTransform
+from rnaglib.dataset_transforms import CDHitComputer, ClusterSplitter, StructureDistanceComputer, RandomSplitter
+from rnaglib.encoders import ListEncoder
+from rnaglib.config.graph_keys import GRAPH_KEYS, TOOL
 
 script_dir = os.path.dirname(os.path.realpath(__file__))
 if __name__ == "__main__":
     sys.path.append(os.path.join(script_dir, '..'))
 
 from exp import RNATrainer
 
-# Setup task
-ta = get_task(root="roots/RNA_CM", task_id="rna_cm")
+# Hyperparameters (to tune)
+nb_layers = 3
+hidden_dim = 128
+learning_rate = 0.001
+batch_size = 8
+epochs = 40
+split = "default"
+rna_fm = False
+representation = "2.5D"
+layer_type = "rgcn"
+output = "tensorboard"
+
+# Experiment name
+exp_name="RNA_CM_"+str(nb_layers)+"layers_lr"+str(learning_rate)+"_"+str(epochs)+"epochs_hiddendim"+str(hidden_dim)+"_"+representation+"_layer_type_"+layer_type
+if rna_fm:
+    exp_name += "rna_fm"
+if split != "default":
+     exp_name += split
 
-ta.dataset.add_representation(GraphRepresentation(framework="pyg"))
-ta.get_split_loaders(batch_size=8)
 
 model_args = {
-        "num_node_features": ta.metadata["num_node_features"],
-        "num_classes": ta.metadata["num_classes"],
-        "graph_level": False,
-        "num_layers": 3,
+    "graph_level": False,
+    "num_layers": nb_layers,
+    "hidden_channels": hidden_dim,
+    "layer_type": layer_type,
 }
 
-model = PygModel(**model_args)
-trainer = RNATrainer(ta, model, epochs=40)
+if rna_fm:
+    model_args["num_node_features"]=644
 
-trainer.train()
+#model_CM = PygModel(**model_args)
+#trainer_CM = RNATrainer(ta, model_CM, rep, exp_name=exp_name, learning_rate=learning_rate, epochs=epochs)
+#trainer_CM.train()
 
 if __name__ == "__main__":
-    pass
+    for seed in [0,1,2]:
+        ta = get_task(root="roots/RNA_CM", task_id="rna_cm")
+        if split=="struc":
+            distance = "USalign"
+        else:
+                distance = "cd_hit"
+
+        if distance not in ta.dataset.distances:
+            if split == 'struc':
+                    ta.dataset = StructureDistanceComputer()(ta.dataset)
+            if split == 'seq':
+                    ta.dataset = CDHitComputer()(ta.dataset)
+        if split == 'rand':
+            ta.splitter = RandomSplitter()
+        elif split=='struc' or split=='seq':
+            ta.splitter = ClusterSplitter(distance_name=distance)
+
+        if rna_fm:
+            rnafm = RNAFMTransform()
+            [rnafm(rna) for rna in ta.dataset]
+            ta.dataset.features_computer.add_feature(feature_names=["rnafm"], custom_encoders={"rnafm": ListEncoder(640)})
+
+        if representation=="2D":
+            edge_map = GRAPH_KEYS["2D_edge_map"][TOOL]
+        elif representation=="simplified_2.5D":
+            edge_map = GRAPH_KEYS["simplified_edge_map"][TOOL]
+        else:
+            edge_map = GRAPH_KEYS["edge_map"][TOOL]
+
+        representation_args = {
+            "framework": "pyg",
+            "edge_map": edge_map,
+        }
+
+        rep = GraphRepresentation(**representation_args)
+        ta.dataset.add_representation(rep)
+        ta.get_split_loaders(batch_size=batch_size, recompute=True)
+        model = PygModel.from_task(ta, **model_args)
+        trainer = RNATrainer(ta, model, rep, exp_name=exp_name+"_seed"+str(seed), learning_rate=learning_rate, epochs=epochs, seed=seed, batch_size=batch_size, output=output)
+        trainer.train()

exp.py

@@ -6,10 +6,11 @@
 
 import torch
 import numpy as np
+from torch.utils.tensorboard import SummaryWriter
 
 class RNATrainer:
     def __init__(self, task, model, rep="pyg", wandb_project="", exp_name="default",
-                 learning_rate=0.001, epochs=100, seed=0, batch_size=8):
+                 learning_rate=0.001, epochs=100, seed=0, batch_size=8, output="wandb", log_dir="runs/"):
         self.task = task
         self.representation = rep
         self.model = model
@@ -20,15 +21,20 @@ def __init__(self, task, model, rep="pyg", wandb_project="", exp_name="default",
         self.training_log = []
         self.seed = seed
         self.batch_size = batch_size
+        self.output = output
+        self.log_dir = log_dir
 
     def setup(self):
         """Initialize wandb and model training"""
-        wandb.init(
-            entity="mlsb",  # Replace with your team name
-            project=self.wandb_project,
-            name=self.exp_name,
-        )
-    
+        if self.output == "tensorboard":
+            self.train_writer = SummaryWriter(log_dir=self.log_dir+self.task.name+"/"+self.exp_name+"/train")
+            self.val_writer = SummaryWriter(log_dir=self.log_dir+self.task.name+"/"+self.exp_name+"/val")
+        else:
+            wandb.init(
+                entity="mlsb",  # Replace with your team name
+                project=self.wandb_project,
+                name=self.exp_name,
+            )
         # Set seeds for reproducibility
         torch.manual_seed(self.seed)  # CPU random number generator
         if torch.cuda.is_available():
@@ -72,33 +78,48 @@ def train(self):
             train_metrics = self.model.evaluate(self.task, split="train")
             val_metrics = self.model.evaluate(self.task, split="val")
 
-            # Log to wandb
+            # Log to wandb or Tensorboard
             metrics = {
                 "epoch": epoch,
                 **{f"train_{k}": v for k, v in train_metrics.items()},
                 **{f"val_{k}": v for k, v in val_metrics.items()}
             }
+            if self.output == "tensorboard":
+                self.train_writer.add_scalar("Loss", train_metrics['loss'], epoch)
+                self.val_writer.add_scalar("Loss", val_metrics['loss'], epoch)
             try:
                 metrics["train_auc"] = train_metrics['auc']
                 metrics["val_auc"] = val_metrics['auc']
+                if self.output == "tensorboard":
+                    self.train_writer.add_scalar("AUC", train_metrics['auc'], epoch)
+                    self.val_writer.add_scalar("AUC", val_metrics['auc'], epoch)
             except:
                 pass
             if self.task.metadata['multi_label']:
                 metrics["train_jaccard"] = train_metrics["jaccard"]
                 metrics["val_jaccard"] = val_metrics["jaccard"]
+                if self.output == "tensorboard":
+                    self.train_writer.add_scalar("Jaccard", train_metrics['jaccard'], epoch)
+                    self.val_writer.add_scalar("Jaccard", val_metrics['jaccard'], epoch)
             else:
                 try:
                     metrics["train_balanced_accuracy"] = train_metrics["balanced_accuracy"]
                     metrics["val_balanced_accuracy"] = val_metrics["balanced_accuracy"]
+                    if self.output == "tensorboard":
+                        self.train_writer.add_scalar("Balanced_acc", train_metrics['balanced_accuracy'], epoch)
+                        self.val_writer.add_scalar("Balanced_acc", val_metrics['balanced_accuracy'], epoch)
                 except:
                     pass
                 try:
                     metrics["train_mcc"] = train_metrics["mcc"]
                     metrics["val_mcc"] = val_metrics["mcc"]
+                    if self.output == "tensorboard":
+                        self.train_writer.add_scalar("MCC", train_metrics['mcc'], epoch)
+                        self.val_writer.add_scalar("MCC", val_metrics['mcc'], epoch)
                 except:
                     pass
-
-            wandb.log(metrics)
+            if self.output == "wandb":
+                wandb.log(metrics)
             self.training_log.append(metrics)
 
             # Print progress
@@ -110,7 +131,11 @@ def train(self):
                 )
 
         self.save_results()
-        wandb.finish()
+        if self.output == "tensorboard":
+            self.train_writer.flush()
+            self.val_writer.flush()
+        else:
+            wandb.finish()
 
     def save_results(self):
         """Save final results and metrics"""