Enhancement/api tools (#168)

* Expose train and predict, train_model returns model

* minor changes for readability

* bugfix for multichtiffdataset"

* try emptying cache after every iter to save memory

* Enhancement for multi-channel support

Author: Gregory Johnson

examples/download_and_train.py

@@ -1,10 +1,11 @@
 import argparse
 import os
 import json
+from pathlib import Path
 
 import quilt3
 import pandas as pd
-from pathlib import Path
+import numpy as np
 
 from fnet.cli.init import save_default_train_options
 
@@ -50,11 +51,17 @@
 
 data_manifest = aics_pipeline["metadata.csv"]()
 
+# THE ROWS OF THE MANIFEST CORRESPOND TO CELLS, WE TRIM DOWN TO UNIQUIE FOVS
+unique_fov_indices = np.unique(data_manifest['FOVId'], return_index=True)[1]
+data_manifest = data_manifest.iloc[unique_fov_indices]
+
 # SELECT THE FIRST N_IMAGES_TO_DOWNLOAD
 data_manifest = data_manifest.iloc[0:n_images_to_download]
 
 image_source_paths = data_manifest["SourceReadPath"]
 
+pdb.set_trace()
+
 image_target_paths = [
     "{}/{}".format(image_save_dir, image_source_path)
     for image_source_path in image_source_paths

fnet/cli/train_model.py

@@ -99,7 +99,7 @@ def main(args: Optional[argparse.Namespace] = None):
 
     if args.path_json and not args.path_json.exists():
         save_default_train_options(args.path_json)
-        return
+        return None
 
     with open(args.path_json, "r") as fi:
         train_options = json.load(fi)
@@ -125,7 +125,7 @@ def main(args: Optional[argparse.Namespace] = None):
 
     if (args.n_iter - model.count_iter) <= 0:
         # Stop if no more iterations needed
-        return
+        return model
 
     # Get patch pair providers
     bpds_train = get_bpds_train(args)
@@ -136,12 +136,14 @@ def main(args: Optional[argparse.Namespace] = None):
         do_save = ((idx_iter + 1) % args.interval_save == 0) or (
             (idx_iter + 1) == args.n_iter
         )
+
         loss_train = model.train_on_batch(*bpds_train.get_batch(args.batch_size))
         loss_val = None
         if do_save and bpds_val is not None:
             loss_val = model.test_on_iterator(
                 [bpds_val.get_batch(args.batch_size) for _ in range(4)]
             )
+
         fnetlogger.add(
             {"num_iter": idx_iter + 1, "loss_train": loss_train, "loss_val": loss_val}
         )

fnet/data/multichtiffdataset.py

@@ -24,15 +24,18 @@ def __init__(
         # if this column is a string assume it is in "[ind_1, ind_2, ..., ind_n]" format
         if isinstance(self.df["channel_signal"][0], str):
             self.df["channel_signal"] = [
-                np.fromstring(ch[1:-1], sep=" ").astype(int)
+                np.fromstring(ch[1:-1], sep=", ").astype(int)
                 for ch in self.df["channel_signal"]
             ]
+        else:
+            self.df["channel_signal"] = [[int(ch)] for ch in self.df["channel_signal"]]
+
+        if isinstance(self.df["channel_target"][0], str):
             self.df["channel_target"] = [
-                np.fromstring(ch[1:-1], sep=" ").astype(int)
+                np.fromstring(ch[1:-1], sep=", ").astype(int)
                 for ch in self.df["channel_target"]
             ]
         else:
-            self.df["channel_signal"] = [[int(ch)] for ch in self.df["channel_signal"]]
             self.df["channel_target"] = [[int(ch)] for ch in self.df["channel_target"]]
 
         assert all(

fnet/fnet_model.py

@@ -232,6 +232,7 @@ def train_on_batch(
             module = torch.nn.DataParallel(self.net, device_ids=self.gpu_ids)
         else:
             module = self.net
+
         self.optimizer.zero_grad()
         y_hat_batch = module(x_batch)
         args = [y_hat_batch, y_batch]

fnet/nn_modules/fnet_nn_3d_params.py

@@ -10,7 +10,7 @@ def __init__(self, depth=4, mult_chan=32, in_channels=1, out_channels=1):
         self.out_channels = out_channels
 
         self.net_recurse = _Net_recurse(
-            n_in_channels=self.in_channels, mult_chan=self.mult_chan, depth=self.depth
+            n_in_channels=self.in_channels, mult_chan=self.mult_chan, depth_parent=self.depth, depth=self.depth
         )
         self.conv_out = torch.nn.Conv3d(
             self.mult_chan, self.out_channels, kernel_size=3, padding=1
@@ -22,7 +22,7 @@ def forward(self, x):
 
 
 class _Net_recurse(torch.nn.Module):
-    def __init__(self, n_in_channels, mult_chan=2, depth=0):
+    def __init__(self, n_in_channels, mult_chan=2, depth_parent=0, depth=0):
         """Class for recursive definition of U-network.p
 
         Parameters
@@ -37,8 +37,14 @@ def __init__(self, n_in_channels, mult_chan=2, depth=0):
 
         """
         super().__init__()
+
         self.depth = depth
-        n_out_channels = n_in_channels * mult_chan
+
+        if self.depth == depth_parent:
+            n_out_channels = mult_chan
+        else:
+            n_out_channels = n_in_channels * mult_chan
+
         self.sub_2conv_more = SubNet2Conv(n_in_channels, n_out_channels)
         if depth > 0:
             self.sub_2conv_less = SubNet2Conv(2 * n_out_channels, n_out_channels)
@@ -52,7 +58,7 @@ def __init__(self, n_in_channels, mult_chan=2, depth=0):
             )
             self.bn1 = torch.nn.BatchNorm3d(n_out_channels)
             self.relu1 = torch.nn.ReLU()
-            self.sub_u = _Net_recurse(n_out_channels, mult_chan=2, depth=(depth - 1))
+            self.sub_u = _Net_recurse(n_out_channels, mult_chan=2, depth_parent=depth_parent, depth=(depth - 1))
 
     def forward(self, x):
         if self.depth == 0:

fnet/tests/data/testlib.py

@@ -70,8 +70,8 @@ def create_multichtiff_data(
             {
                 "dummy_id": idx,
                 "path_tiff": path_x,
-                "channel_signal": np.arange(0, n_ch_in),
-                "channel_target": np.arange(0, n_ch_out) + n_ch_in,
+                "channel_signal": list(np.arange(0, n_ch_in)),
+                "channel_target": list(np.arange(0, n_ch_out) + n_ch_in),
             }
         )