From cbe5aab52e8fe146e463503988590797e70267af Mon Sep 17 00:00:00 2001 From: anna-grim Date: Thu, 2 Jul 2026 19:24:15 +0000 Subject: [PATCH 1/3] refactor: train --- .../machine_learning/train.py | 53 +++++++------------ src/neuron_proofreader/utils/ml_util.py | 33 ++++++++++++ src/neuron_proofreader/utils/swc_util.py | 2 + 3 files changed, 53 insertions(+), 35 deletions(-) diff --git a/src/neuron_proofreader/machine_learning/train.py b/src/neuron_proofreader/machine_learning/train.py index 8fc5ec1..83e19e1 100644 --- a/src/neuron_proofreader/machine_learning/train.py +++ b/src/neuron_proofreader/machine_learning/train.py @@ -191,7 +191,7 @@ def train_step(self, dataloader, epoch): # Write stats to tensorboard stats = metrics.compute() - self.writer.add_scalars("train", stats, epoch) + self.update_tensorboard(stats, epoch, "train_") return stats def validate_step(self, dataloader, epoch): @@ -243,7 +243,7 @@ def validate_step(self, dataloader, epoch): # Write stats to tensorboard stats = metrics.compute() - self.writer.add_scalars("train", stats, epoch) + self.update_tensorboard(stats, epoch, "val_") return stats def forward_pass(self, x, y): @@ -272,39 +272,6 @@ def forward_pass(self, x, y): return y, hat_y, loss # --- Helpers --- - @staticmethod - def compute_stats(y, hat_y): - """ - Computes F1 score, precision, and recall for each sample in a batch. - - Parameters - ---------- - y : torch.Tensor - Ground truth labels with shape (B, 1). - hat_y : torch.Tensor - Predicted labels with shape (B, 1). - - Returns - ------- - stats : Dict[str, float] - Dictionary of metric names to values. - """ - # Binarize predictions - hat_y = (hat_y > 0).int() - y = y.int() - - # Compute stats - tp = (hat_y * y).sum() - fp = (hat_y * (1 - y)).sum() - fn = ((1 - hat_y) * y).sum() - - precision = tp / (tp + fp + 1e-8) - recall = tp / (tp + fn + 1e-8) - f1 = 2 * precision * recall / (precision + recall + 1e-8) - - stats = {"f1": f1, "precision": precision, "recall": recall} - return stats - @staticmethod def report_stats(stats, is_train=True): """ @@ -405,6 +372,22 @@ def save_model(self, epoch): path = os.path.join(self.log_dir, filename) torch.save(self.model.state_dict(), path) + def update_tensorboard(self, stats, epoch, prefix): + """ + Logs scalar statistics to TensorBoard. + + Parameters + ---------- + stats : Dict[str, float] + Dictionary of metric names to lists of values. + epoch : int + Current training epoch. + prefix : str + Prefix to prepend to each metric name when logging. + """ + for key, value in stats.items(): + self.writer.add_scalar(prefix + key, stats[key], epoch) + class DistributedTrainer(Trainer): """ diff --git a/src/neuron_proofreader/utils/ml_util.py b/src/neuron_proofreader/utils/ml_util.py index 5b67be1..d557796 100644 --- a/src/neuron_proofreader/utils/ml_util.py +++ b/src/neuron_proofreader/utils/ml_util.py @@ -213,6 +213,39 @@ def move(self, v, device): # --- Miscellaneous --- +def find_max_batch_size( + model, input_shape, optimizer_cls, device="cuda", start=1, max_bs=32 +): + model.to(device) + lo, hi = start, max_bs + best = 0 + while lo <= hi: + bs = (lo + hi) // 2 + try: + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + x = torch.randn(bs, *input_shape, device=device) + y = torch.randn(bs, 1, device=device) + opt = optimizer_cls(model.parameters()) + scaler = torch.cuda.amp.GradScaler() + + with torch.autocast(device_type="cuda", dtype=torch.float16): + out = model(x) + loss = F.mse_loss(out, y) + scaler.scale(loss).backward() + scaler.step(opt) + scaler.update() + + best = bs + lo = bs + 1 + del x, y, out, loss, opt + except torch.cuda.OutOfMemoryError: + hi = bs - 1 + finally: + torch.cuda.empty_cache() + return best + + def load_model(model, model_path, device="cuda"): """ Loads a PyTorch model checkpoint, moves the model to the speficied device, diff --git a/src/neuron_proofreader/utils/swc_util.py b/src/neuron_proofreader/utils/swc_util.py index f301117..5be01bb 100644 --- a/src/neuron_proofreader/utils/swc_util.py +++ b/src/neuron_proofreader/utils/swc_util.py @@ -292,6 +292,8 @@ def read_from_cloud(self, path): # List paths swc_paths = util.list_cloud_paths(path, ".swc") zip_paths = util.list_cloud_paths(path, ".zip") + print("path:", path) + print("zip paths:", zip_paths) # Call reader if swc_paths: From f8ece1ec261a6c88af73cc5781d28f0e367b1da7 Mon Sep 17 00:00:00 2001 From: anna-grim Date: Thu, 2 Jul 2026 19:25:23 +0000 Subject: [PATCH 2/3] refactor: removed prints --- src/neuron_proofreader/utils/swc_util.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/neuron_proofreader/utils/swc_util.py b/src/neuron_proofreader/utils/swc_util.py index 5be01bb..f301117 100644 --- a/src/neuron_proofreader/utils/swc_util.py +++ b/src/neuron_proofreader/utils/swc_util.py @@ -292,8 +292,6 @@ def read_from_cloud(self, path): # List paths swc_paths = util.list_cloud_paths(path, ".swc") zip_paths = util.list_cloud_paths(path, ".zip") - print("path:", path) - print("zip paths:", zip_paths) # Call reader if swc_paths: From 154f303b5761ea5843af71d1b353164302cb2f7a Mon Sep 17 00:00:00 2001 From: anna-grim Date: Thu, 2 Jul 2026 19:50:31 +0000 Subject: [PATCH 3/3] refactor: gcs client attr in swc read --- src/neuron_proofreader/utils/ml_util.py | 3 ++- src/neuron_proofreader/utils/swc_util.py | 10 +++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/neuron_proofreader/utils/ml_util.py b/src/neuron_proofreader/utils/ml_util.py index d557796..fa42897 100644 --- a/src/neuron_proofreader/utils/ml_util.py +++ b/src/neuron_proofreader/utils/ml_util.py @@ -9,9 +9,9 @@ """ -import numpy as np import torch import torch.nn as nn +import torch.nn.functional as F # --- Architectures --- @@ -232,6 +232,7 @@ def find_max_batch_size( with torch.autocast(device_type="cuda", dtype=torch.float16): out = model(x) loss = F.mse_loss(out, y) + scaler.scale(loss).backward() scaler.step(opt) scaler.update() diff --git a/src/neuron_proofreader/utils/swc_util.py b/src/neuron_proofreader/utils/swc_util.py index f301117..f0ffa73 100644 --- a/src/neuron_proofreader/utils/swc_util.py +++ b/src/neuron_proofreader/utils/swc_util.py @@ -47,6 +47,8 @@ class Reader: archive, and (3) local directory of ZIP archives. """ + gcs_client = None + def __init__( self, anisotropy=(1.0, 1.0, 1.0), min_swc_pts=1, verbose=True ): @@ -67,6 +69,12 @@ def __init__( self.min_swc_pts = min_swc_pts self.verbose = verbose + @classmethod + def _get_gcs_client(cls): + if cls._gcs_client is None: + cls._gcs_client = storage.Client() + return cls._gcs_client + # --- Read Data --- def __call__(self, swc_pointer): """ @@ -319,7 +327,7 @@ def read_gcs_swc(self, path): """ # Initialize cloud reader bucket_name, key = util.parse_cloud_path(path) - bucket = storage.Client().bucket(bucket_name) + bucket = self._get_gcs_client().bucket(bucket_name) blob = bucket.blob(key) # Parse swc contents