Enhancement/api tools (#164)

* Expose train and predict, train_model returns model

* minor changes for readability

* bugfix for multichtiffdataset"

Author: Gregory Johnson

fnet/__init__.py

@@ -1,6 +1,10 @@
 from fnet import models
 from fnet.fnetlogger import FnetLogger
 
+# Clean these up later - GRJ 2020-02-04
+from fnet.cli.train_model import train_model as train
+from fnet.cli.predict import main as predict
+
 __author__ = "Gregory R. Johnson"
 __email__ = "[email protected]"
 __version__ = "0.2.0"

fnet/cli/main.py

@@ -28,10 +28,12 @@ def main() -> None:
     init.add_parser_arguments(parser_init)
     train_model.add_parser_arguments(parser_train)
     predict.add_parser_arguments(parser_predict)
+
     parser_init.set_defaults(func=init.main)
     parser_train.set_defaults(func=train_model.main)
     parser_predict.set_defaults(func=predict.main)
     args = parser.parse_args()
+
     # Remove 'func' from args so it is not passed to target script
     func = args.func
     delattr(args, "func")

fnet/cli/train_model.py

@@ -86,26 +86,32 @@ def add_parser_arguments(parser) -> None:
     parser.add_argument("--gpu_ids", nargs="+", default=[0], type=int, help="gpu_id(s)")
 
 
-def main(args: Optional[argparse.Namespace] = None) -> None:
+def main(args: Optional[argparse.Namespace] = None):
     """Trains a model."""
     time_start = time.time()
+
     if args is None:
         parser = argparse.ArgumentParser()
         add_parser_arguments(parser)
         args = parser.parse_args()
-    if args.json and not args.json.exists():
-        save_default_train_options(args.json)
+
+    args.path_json = Path(args.json)
+
+    if args.path_json and not args.path_json.exists():
+        save_default_train_options(args.path_json)
         return
-    with open(args.json, "r") as fi:
+
+    with open(args.path_json, "r") as fi:
         train_options = json.load(fi)
+
     args.__dict__.update(train_options)
     add_logging_file_handler(Path(args.path_save_dir, "train_model.log"))
     logger.info(f"Started training at: {datetime.datetime.now()}")
 
     set_seeds(args.seed)
     log_training_options(vars(args))
     path_model = os.path.join(args.path_save_dir, "model.p")
-    model = fnet.models.load_or_init_model(path_model, args.json)
+    model = fnet.models.load_or_init_model(path_model, args.path_json)
     init_cuda(args.gpu_ids[0])
     model.to_gpu(args.gpu_ids)
     logger.info(model)
@@ -124,6 +130,8 @@ def main(args: Optional[argparse.Namespace] = None) -> None:
     # Get patch pair providers
     bpds_train = get_bpds_train(args)
     bpds_val = get_bpds_val(args)
+
+    # MAIN LOOP
     for idx_iter in range(model.count_iter, args.n_iter):
         do_save = ((idx_iter + 1) % args.interval_save == 0) or (
             (idx_iter + 1) == args.n_iter
@@ -164,6 +172,8 @@ def main(args: Optional[argparse.Namespace] = None) -> None:
                 path_save=os.path.join(args.path_save_dir, "loss_curves.png"),
             )
 
+    return model
+
 
 def train_model(
     batch_size: int = 28,
@@ -182,8 +192,9 @@ def train_model(
     seed: Optional[int] = None,
     json: Optional[str] = None,
     gpu_ids: Optional[List[int]] = None,
-) -> None:
+):
     """Python API for training."""
+
     bpds_kwargs = bpds_kwargs or {
         "buffer_size": 16,
         "buffer_switch_interval": 2800,  # every 100 updates
@@ -201,7 +212,8 @@ def train_model(
     }
     iter_checkpoint = iter_checkpoint or []
     gpu_ids = gpu_ids or [0]
-    json = json or str(Path(path_save_dir, "train_options.json"))
+
+    json = json or f"{path_save_dir}train_options.json"
 
     pnames, _, _, locs = inspect.getargvalues(inspect.currentframe())
     train_options = {k: locs[k] for k in pnames}
@@ -214,10 +226,11 @@ def train_model(
         path_json.parent.mkdir(parents=True)
 
     json = globals()["json"]  # retrieve global module
-    with path_json.open("w") as fo:
-        json.dump(train_options, fo, indent=4, sort_keys=True)
+    with path_json.open("w") as f:
+        json.dump(train_options, f, indent=4, sort_keys=True)
         logger.info(f"Saved: {path_json}")
 
     args = argparse.Namespace()
     args.__dict__.update(train_options)
-    main(args)
+
+    return main(args)

fnet/data/multichtiffdataset.py

@@ -7,10 +7,8 @@
 
 
 class MultiChTiffDataset(FnetDataset):
-    """Dataset for multi-channel tiff files.
-
-    Currently assumes that images are loaded in STCZYX format
-
+    """
+    Dataset for multi-channel tiff files.
     """
 
     def __init__(
@@ -43,6 +41,16 @@ def __init__(
         )
 
     def __getitem__(self, index):
+        """
+        Parameters
+        ----------
+        index: integer
+
+        Returns
+        -------
+        C by <spatial dimensions> torch.Tensor
+        """
+
         element = self.df.iloc[index, :]
         has_target = not np.any(np.isnan(element["channel_target"]))
 
@@ -67,7 +75,7 @@ def __getitem__(self, index):
         im_out = [torch.from_numpy(im.astype(float)).float() for im in im_out]
 
         # unsqueeze to make the first dimension be the channel dimension
-        im_out = [torch.unsqueeze(im, 0) for im in im_out]
+        # im_out = [torch.unsqueeze(im, 0) for im in im_out]
 
         return tuple(im_out)
 

fnet/tests/test_multichtiffdataset.py

@@ -28,5 +28,5 @@ def test_MultiTiffDataset(tmp_path, n_ch_in, n_ch_out, dims_zyx):
     len_data = 2
     assert len(data) == len_data
 
-    assert tuple(data[0].shape) == (1,) + (n_ch_in,) + dims_zyx
-    assert tuple(data[1].shape) == (1,) + (n_ch_out,) + dims_zyx
+    assert tuple(data[0].shape) == (n_ch_in,) + dims_zyx
+    assert tuple(data[1].shape) == (n_ch_out,) + dims_zyx