Use logging

fnet/cli/init.py

@@ -1,11 +1,16 @@
 from typing import Optional
+from pathlib import Path
 import argparse
 import json
+import logging
 import os
 import shutil
 import sys
 
 
+logger = logging.getLogger(__name__)
+
+
 def save_example_scripts(path_save_dir: str) -> None:
     """Save example training and prediction scripts.
 
@@ -24,10 +29,10 @@ def save_example_scripts(path_save_dir: str) -> None:
         path_src = os.path.join(path_examples_dir, fname)
         path_dst = os.path.join(path_save_dir, fname)
         if os.path.exists(path_dst):
-            print('Example script already exists:', path_dst)
+            logger.info(f'Example script already exists: {path_dst}')
             continue
         shutil.copy(path_src, path_dst)
-        print('Saved:', path_dst)
+        logger.info(f'Saved: {path_dst}')
 
 
 def save_default_train_options(path_save: str) -> None:
@@ -39,12 +44,11 @@ def save_default_train_options(path_save: str) -> None:
         Save path for default training options json.
 
     """
-    if os.path.exists(path_save):
-        print('Training options file already exists:', path_save)
+    path_save = Path(path_save)
+    if path_save.exists():
+        logger.info(f'Training options file already exists: {path_save}')
         return
-    dirname = os.path.dirname(path_save)
-    if not os.path.exists(dirname):
-        os.makedirs(dirname)
+    path_save.parent.mkdir(parents=True, exist_ok=True)
     train_options = {
         'batch_size': 28,
         'bpds_kwargs': {
@@ -69,12 +73,12 @@ def save_default_train_options(path_save: str) -> None:
         'interval_save': 1000,
         'iter_checkpoint': [],
         'n_iter': 250000,
-        'path_save_dir': dirname,
+        'path_save_dir': str(path_save.parent),
         'seed': None,
     }
-    with open(path_save, 'w') as fo:
+    with path_save.open('w') as fo:
         json.dump(train_options, fo, indent=4, sort_keys=True)
-        print('Saved:', path_save)
+        logger.info(f'Saved: {path_save}')
 
 
 def add_parser_arguments(parser: argparse.ArgumentParser) -> None:

fnet/cli/main.py

@@ -8,10 +8,12 @@
 from fnet.cli import init
 from fnet.cli import predict
 from fnet.cli import train_model
+from fnet.utils.general_utils import init_fnet_logging
 
 
 def main() -> None:
     """Main function for command-line 'fnet' command."""
+    init_fnet_logging()
     parser = argparse.ArgumentParser(prog='fnet')
     subparser = parser.add_subparsers(title='command')
     parser_init = subparser.add_parser(

fnet/cli/predict.py

@@ -4,6 +4,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple
 import argparse
 import json
+import logging
 import os
 
 import numpy as np
@@ -19,6 +20,9 @@
 from fnet.utils.general_utils import str_to_object
 
 
+logger = logging.getLogger(__name__)
+
+
 def get_dataset(args: argparse.Namespace) -> torch.utils.data.Dataset:
     """Returns dataset.
 
@@ -126,10 +130,10 @@ def save_tif(fname: str, ar: np.ndarray, path_root: str) -> str:
     path_tif_dir = os.path.join(path_root, 'tifs')
     if not os.path.exists(path_tif_dir):
         os.makedirs(path_tif_dir)
-        print('Created:', path_tif_dir)
+        logger.info(f'Created: {path_tif_dir}')
     path_save = os.path.join(path_tif_dir, fname)
     tifffile.imsave(path_save, ar, compress=2)
-    print('Saved:', path_save)
+    logger.info('Saved: {path_save}')
     return os.path.relpath(path_save, path_root)
 
 
@@ -182,9 +186,9 @@ def save_predictions_csv(
     dirname = os.path.dirname(path_csv)
     if not os.path.exists(dirname):
         os.makedirs(dirname)
-        print('Created:', dirname)
+        logger.info('Created: {dirname}')
     retry_if_oserror(df.to_csv)(path_csv)
-    print('Saved:', path_csv)
+    logger.info('Saved: {path_csv}')
 
 
 def save_args_as_json(path_save_dir: str, args: argparse.Namespace) -> None:
@@ -213,7 +217,7 @@ def save_args_as_json(path_save_dir: str, args: argparse.Namespace) -> None:
         )
     with open(path_json, 'w') as fo:
         json.dump(vars(args), fo, indent=4, sort_keys=True)
-        print('Saved:', path_json)
+    logger.info('Saved: {path_json}')
 
 
 def add_parser_arguments(parser) -> None:
@@ -244,7 +248,7 @@ def main(args: Optional[argparse.Namespace] = None) -> None:
     model = None
     indices = get_indices(args, dataset)
     for count, idx in enumerate(indices, 1):
-        print(f'Processing: {idx:3d} ({count}/{len(indices)})')
+        logger.info(f'Processing: {idx:3d} ({count}/{len(indices)})')
         entry = {}
         entry['index'] = idx
         signal, target = item_from_dataset(dataset, idx)
@@ -261,7 +265,7 @@ def main(args: Optional[argparse.Namespace] = None) -> None:
                 model_def = parse_model(path_model_dir)
                 model = load_model(model_def['path'], no_optim=True)
                 model.to_gpu(args.gpu_ids)
-                print('Loaded model:', model_def['name'])
+                logger.info(f'Loaded model: {model_def["name"]}')
             prediction = model.predict_piecewise(
                 signal,
                 tta=('no_tta' not in model_def['options']),

fnet/cli/train_model.py

@@ -1,27 +1,40 @@
 """Trains a model."""
 
 
-from typing import Callable, Optional
+from pathlib import Path
+from typing import Callable, Dict, Optional
 import argparse
 import copy
 import datetime
 import json
 import logging
 import os
 import pprint
-import sys
 import time
 
 import numpy as np
 import torch
 import torch.utils.data
 
 from fnet.cli.init import save_default_train_options
+from fnet.utils.general_utils import add_logging_file_handler
 from fnet.utils.general_utils import str_to_object
 import fnet
 import fnet.utils.viz_utils as vu
 
 
+logger = logging.getLogger(__name__)
+
+
+def log_training_options(options: Dict) -> None:
+    """Logs training options."""
+    for line in (
+            ['*** Training options ***']
+            + pprint.pformat(options).split(os.linesep)
+    ):
+        logger.info(line)
+
+
 def set_seeds(seed: Optional[int]) -> None:
     """Sets random seeds"""
     if seed is None:
@@ -35,8 +48,11 @@ def init_cuda(gpu: int) -> None:
     """Initialize Pytorch CUDA state."""
     if gpu < 0:
         return
-    torch.cuda.set_device(gpu)
-    torch.cuda.init()
+    try:
+        torch.cuda.set_device(gpu)
+        torch.cuda.init()
+    except RuntimeError:
+        logger.exception('Failed to init CUDA')
 
 
 def get_dataloader_train(
@@ -74,30 +90,6 @@ def get_dataloader_val(
     return dataloader
 
 
-def init_logger(path_save: str) -> None:
-    """Initialize training logger.
-
-    Parameters
-    ----------
-    path_save
-        Location to save training log.
-
-    """
-    dirname = os.path.dirname(path_save)
-    if not os.path.exists(dirname):
-        os.makedirs(dirname)
-    logging.basicConfig(
-        handlers=[
-            logging.StreamHandler(sys.stdout),
-            logging.FileHandler(path_save, mode='a'),
-        ]
-    )
-    logger = logging.getLogger(__name__)
-    logger.setLevel(logging.DEBUG)
-    logger.info('Started training at: %s', datetime.datetime.now())
-    return logger
-
-
 def add_parser_arguments(parser) -> None:
     """Add training script arguments to parser."""
     parser.add_argument('json', help='json with training options')
@@ -117,15 +109,11 @@ def main(args: Optional[argparse.Namespace] = None) -> None:
     with open(args.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('Started training at: %s', datetime.datetime.now())
 
-    logger = init_logger(path_save=os.path.join(args.path_save_dir, 'run.log'))
     set_seeds(args.seed)
-    logger.info(
-        os.linesep.join(
-            ['*** Training options ***', pprint.pformat(vars(args))]
-        )
-    )
-
+    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)
     init_cuda(args.gpu_ids[0])

fnet/data/bufferedpatchdataset.py

@@ -1,8 +1,14 @@
-from fnet.data.fnetdataset import FnetDataset
+import logging
+
 from tqdm import tqdm
 import numpy as np
 import torch
 
+from fnet.data.fnetdataset import FnetDataset
+
+
+logger = logging.getLogger(__name__)
+
 
 class BufferedPatchDataset(FnetDataset):
     """Dataset that provides chunks/patchs from another dataset."""
@@ -54,7 +60,7 @@ def __getitem__(self, index):
             self.counter % self.buffer_switch_frequency == 0
         ):
             if self.verbose:
-                print("Inserting new item into buffer")
+                logger.info("Inserting new item into buffer")
             self.insert_new_element_into_buffer()
         return self.get_random_patch()
 
@@ -81,7 +87,7 @@ def insert_new_element_into_buffer(self):
         self.buffer.append(self.dataset[new_datum_index])
 
         if self.verbose:
-            print("Added item {0}".format(new_datum_index))
+            logger.info(f"Added item {new_datum_index}")
 
     def get_random_patch(self):
         buffer_index = np.random.randint(len(self.buffer))

fnet/fnet_ensemble.py

@@ -1,10 +1,16 @@
-from fnet.fnet_model import Model
-from fnet.utils.general_utils import str_to_class
 from typing import Union, List
-import numpy as np
+import logging
 import os
+
+import numpy as np
 import torch
 
+from fnet.fnet_model import Model
+from fnet.utils.general_utils import str_to_class
+
+
+logger = logging.info(__name__)
+
 
 def _load_model(path_model: str) -> Model:
     """Load saved model from path."""
@@ -89,7 +95,6 @@ def predict(
         for path_model in self.paths_model:
             model = _load_model(path_model)
             model.to_gpu(self.gpu_ids)
-            print(model)
             y_hat = model.predict(x=x, tta=tta)
             if y_hat_mean is None:
                 y_hat_mean = torch.zeros(*y_hat.size())
@@ -116,7 +121,7 @@ def save(self, path_save: str):
         if not os.path.exists(dirname):
             os.makedirs(dirname)
         torch.save(state, path_save)
-        print('Ensemble model saved to:', path_save)
+        logger.info(f'Ensemble model saved to: {path_save}')
 
     # Override
     def load_state(self, state: dict, no_optim: bool = False):

fnet/fnet_model.py

@@ -2,6 +2,7 @@
 
 
 from typing import Union, Iterator, Optional, Tuple
+import logging
 import math
 import os
 
@@ -15,6 +16,9 @@
 from fnet.utils.model_utils import move_optim
 
 
+logger = logging.getLogger(__name__)
+
+
 def _weights_init(m):
     classname = m.__class__.__name__
     if classname.startswith('Conv'):
@@ -185,7 +189,7 @@ def save(self, path_save: str):
         dirname = os.path.dirname(path_save)
         if not os.path.exists(dirname):
             os.makedirs(dirname)
-            print('Created:', dirname)
+            logger.info('Created: {dirname}')
         curr_gpu_ids = self.gpu_ids
         self.to_gpu(-1)
         retry_if_oserror(torch.save)(self.get_state(), path_save)