Refactor train routines

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()

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