refactor: split datasets

src/neuron_proofreader/configs.py

@@ -46,6 +46,15 @@ def save(self, output_dir):
         path = os.path.join(output_dir, f"{self.name}.json")
         util.write_json(path, self.to_dict())
 
+    def __repr__(self):
+        fields = self.to_dict()
+        width = max(len(k) for k in fields)
+        lines = "\n".join(
+            f"  {k:<{width}} = {v!r}"
+            for k, v in fields.items()
+        )
+        return f"{self.__class__.__name__}(\n{lines}\n)"
+
 
 @dataclass
 class GraphConfig(Config):
@@ -124,16 +133,20 @@ class ProposalsConfig(Config):
     allow_nonleaf_proposals : bool
         Indication of whether to generate proposals between leaf and nodes
         with degree 2.
-    proposals_per_leaf : int
+    initial_search_radius : float
+        Initial search radius for generating proposals.
+    max_proposals_per_leaf : int
         Maximum number of proposals generated at leaf nodes.
     trim_endpoints_bool : bool
-        Indication of whether trim endpoints of isolated leaf-to-leaf
-        proposals.
+        True if endpoints of isolated leaf-to-leaf proposals should be 
+        trimmed.
     """
 
     allow_nonleaf_proposals: bool = False
+    initial_search_radius: float = 25
+    max_attempts: int = 2
     max_proposals_per_leaf: int = 3
     min_size_with_proposals: float = 0
-    trim_endpoints_bool: bool = True
-    search_radius: float = 25
+    name: str = "proposals_config"
     search_scaling_factor: float = 1.5
+    trim_endpoints_bool: bool = True

src/neuron_proofreader/machine_learning/train.py

@@ -117,7 +117,7 @@ def __init__(
         self.writer = SummaryWriter(log_dir=log_dir)
 
     # --- Core Routines ---
-    def run(self, train_dataloader, val_dataloader):
+    def __call__(self, train_dataloader, val_dataloader):
         """
         Runs the full training and validation loop.
 
@@ -128,9 +128,7 @@ def run(self, train_dataloader, val_dataloader):
         val_dataloader : torch.utils.data.Dataset
             Dataloader used for validation.
         """
-        exp_name = os.path.basename(os.path.normpath(self.log_dir))
-        val_dataloader.dataset.save_val_summary(self.log_dir)
-        print("\nExperiment:", exp_name)
+        print("\nExperiment:", os.path.basename(self.log_dir))
         for epoch in range(self.max_epochs):
             # Train-Validate
             print("\nEpoch", epoch)
@@ -181,11 +179,8 @@ def train_step(self, dataloader, epoch):
             self.scaler.step(self.optimizer)
             self.scaler.update()
 
-            # Compute metrics
-            hat_y = hat_y > 0
-            metrics.update(hat_y, y, loss)
-
-            # Update progress bar
+            # Updates
+            metrics.update(hat_y > 0, y, loss)
             if self.verbose:
                 pbar.update(1)
 
@@ -229,15 +224,12 @@ def validate_step(self, dataloader, epoch):
             with torch.inference_mode():
                 y, hat_y, loss = self.forward_pass(x, y)
 
-            # Compute metrics
-            hat_y = hat_y > 0
-            metrics.update(hat_y, y, loss)
-
             # Save MIPs of mistakes
             self._save_mistake_mips(x, y, hat_y, idx_offset)
             idx_offset += len(y)
 
-            # Update progress bar
+            # Updates
+            metrics.update(hat_y > 0, y, loss)
             if self.verbose:
                 pbar.update(1)
 
@@ -264,8 +256,8 @@ def forward_pass(self, x, y):
         loss : torch.Tensor
             Computed loss value.
         """
-        x = x.to(self.device, non_blocking=True)
-        y = y.to(self.device, non_blocking=True)
+        x = x.to(self.device)
+        y = y.to(self.device)
         with torch.autocast(device_type="cuda", dtype=torch.float16):
             hat_y = self.model(x)
             loss = self.criterion(hat_y, y)