Makedirs helper function

Author: swansonk14

chemprop/models/__init__.py

@@ -1 +1 @@
-from .model import build_model
+from .model import build_model, MoleculeModel

chemprop/parsing.py

@@ -5,6 +5,8 @@
 
 import torch
 
+from chemprop.utils import makedirs
+
 
 def add_predict_args(parser: ArgumentParser):
     """
@@ -202,9 +204,7 @@ def modify_predict_args(args: Namespace):
     del args.no_cuda
 
     # Create directory for preds path
-    preds_dir = os.path.dirname(args.preds_path)
-    if preds_dir != '':
-        os.makedirs(preds_dir, exist_ok=True)
+    makedirs(args.preds_path, isfile=True)
 
 
 def parse_predict_args() -> Namespace:
@@ -235,7 +235,7 @@ def modify_train_args(args: Namespace):
     assert args.dataset_type is not None
 
     if args.save_dir is not None:
-        os.makedirs(args.save_dir, exist_ok=True)
+        makedirs(args.save_dir)
     else:
         temp_dir = TemporaryDirectory()
         args.save_dir = temp_dir.name

chemprop/train/cross_validate.py

@@ -5,8 +5,9 @@
 
 import numpy as np
 
-from chemprop.data.utils import get_task_names
 from .run_training import run_training
+from chemprop.data.utils import get_task_names
+from chemprop.utils import makedirs
 
 
 def cross_validate(args: Namespace, logger: Logger = None) -> Tuple[float, float]:
@@ -24,7 +25,7 @@ def cross_validate(args: Namespace, logger: Logger = None) -> Tuple[float, float
         info(f'Fold {fold_num}')
         args.seed = init_seed + fold_num
         args.save_dir = os.path.join(save_dir, f'fold_{fold_num}')
-        os.makedirs(args.save_dir, exist_ok=True)
+        makedirs(args.save_dir)
         model_scores = run_training(args, logger)
         all_scores.append(model_scores)
     all_scores = np.array(all_scores)

chemprop/train/run_training.py

@@ -20,7 +20,7 @@
 from chemprop.models import build_model
 from chemprop.nn_utils import param_count
 from chemprop.utils import build_optimizer, build_lr_scheduler, get_loss_func, get_metric_func, load_checkpoint,\
-    save_checkpoint
+    makedirs, save_checkpoint
 
 
 def run_training(args: Namespace, logger: Logger = None) -> List[float]:
@@ -140,7 +140,7 @@ def run_training(args: Namespace, logger: Logger = None) -> List[float]:
     for model_idx in range(args.ensemble_size):
         # Tensorboard writer
         save_dir = os.path.join(args.save_dir, f'model_{model_idx}')
-        os.makedirs(save_dir, exist_ok=True)
+        makedirs(save_dir)
         writer = SummaryWriter(log_dir=save_dir)
 
         # Load/build model

chemprop/utils.py

@@ -12,19 +12,35 @@
 from torch.optim.lr_scheduler import _LRScheduler
 
 from chemprop.data import StandardScaler
-from chemprop.models import build_model
+from chemprop.models import build_model, MoleculeModel
 from chemprop.nn_utils import NoamLR
 
 
+def makedirs(path: str, isfile: bool = False):
+    """
+    Creates a directory given a path to either a directory or file.
+
+    If a directory is provided, creates that directory. If a file is provided (i.e. isfiled == True),
+    creates the parent directory for that file.
+
+    :param path: Path to a directory or file.
+    :param isfile: Whether the provided path is a directory or file.
+    """
+    if isfile:
+        path = os.path.dirname(path)
+    if path != '':
+        os.makedirs(path, exist_ok=True)
+
+
 def save_checkpoint(path: str,
-                    model: nn.Module,
+                    model: MoleculeModel,
                     scaler: StandardScaler = None,
                     features_scaler: StandardScaler = None,
                     args: Namespace = None):
     """
     Saves a model checkpoint.
 
-    :param model: A PyTorch model.
+    :param model: A MoleculeModel.
     :param scaler: A StandardScaler fitted on the data.
     :param features_scaler: A StandardScaler fitted on the features.
     :param args: Arguments namespace.
@@ -48,15 +64,15 @@ def save_checkpoint(path: str,
 def load_checkpoint(path: str,
                     current_args: Namespace = None,
                     cuda: bool = False,
-                    logger: logging.Logger = None) -> nn.Module:
+                    logger: logging.Logger = None) -> MoleculeModel:
     """
     Loads a model checkpoint.
 
     :param path: Path where checkpoint is saved.
     :param current_args: The current arguments. Replaces the arguments loaded from the checkpoint if provided.
     :param cuda: Whether to move model to cuda.
     :param logger: A logger.
-    :return: The loaded model.
+    :return: The loaded MoleculeModel.
     """
     debug = logger.debug if logger is not None else print
 
@@ -275,8 +291,7 @@ def create_logger(name: str, save_dir: str = None, quiet: bool = False) -> loggi
     logger.addHandler(ch)
 
     if save_dir is not None:
-        if save_dir != '':
-            os.makedirs(save_dir, exist_ok=True)
+        makedirs(save_dir)
 
         fh_v = logging.FileHandler(os.path.join(save_dir, 'verbose.log'))
         fh_v.setLevel(logging.DEBUG)

hyperparameter_optimization.py

@@ -13,7 +13,7 @@
 from chemprop.nn_utils import param_count
 from chemprop.parsing import add_train_args, modify_train_args
 from chemprop.train import cross_validate
-from chemprop.utils import create_logger
+from chemprop.utils import create_logger, makedirs
 
 
 SPACE = {
@@ -86,9 +86,7 @@ def objective(hyperparams: Dict[str, Union[int, float]]) -> float:
     logger.info(f'{best_result["mean_score"]} +/- {best_result["std_score"]} {args.metric}')
 
     # Save best hyperparameter settings as JSON config file
-    config_save_dir = os.path.dirname(args.config_save_path)
-    if config_save_dir != '':
-        os.makedirs(config_save_dir, exist_ok=True)
+    makedirs(args.config_save_path, isfile=True)
 
     with open(args.config_save_path, 'w') as f:
         json.dump(best_result['hyperparams'], f, indent=4, sort_keys=True)

scripts/save_features.py

@@ -16,6 +16,7 @@
 
 from chemprop.data.utils import get_data
 from chemprop.features import get_features_func
+from chemprop.utils import makedirs
 
 
 def load_temp(temp_dir: str) -> Tuple[List[List[float]], int]:
@@ -62,6 +63,9 @@ def save_features(args: Namespace):
 
     :param args: Arguments.
     """
+    # Create directory for save_path
+    makedirs(args.save_path, isfile=True)
+
     # Get data and features function
     data = get_data(path=args.data_path, max_data_size=None)
     features_func = get_features_func(args.features_generator)
@@ -81,7 +85,7 @@ def save_features(args: Namespace):
             features, temp_num = load_temp(temp_save_dir)
 
     if not os.path.exists(temp_save_dir):
-        os.makedirs(temp_save_dir)
+        makedirs(temp_save_dir)
         features, temp_num = [], 0
 
     # Build features map function
@@ -134,8 +138,4 @@ def save_features(args: Namespace):
                         help='Whether to run in parallel rather than sequentially (warning: doesn\'t always work')
     args = parser.parse_args()
 
-    dirname = os.path.dirname(args.save_path)
-    if dirname != '':
-        os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
-
     save_features(args)

test/test.py

@@ -13,7 +13,7 @@
 from chemprop.features import clear_cache
 from chemprop.parsing import add_train_args, modify_train_args, add_predict_args, modify_predict_args
 from chemprop.train import cross_validate, make_predictions
-from chemprop.utils import create_logger
+from chemprop.utils import create_logger, makedirs
 
 from hyperparameter_optimization import grid_search
 
@@ -91,9 +91,7 @@ def test_save_features(self):
             args.save_path = NamedTemporaryFile().name
             args.features_generator = 'morgan_count'
 
-            dirname = os.path.dirname(args.save_path)
-            if dirname != '':
-                os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
+            makedirs(args.save_path, isfile=True)
 
             save_features(args)
             os.unlink(args.save_path)

web/web.py

@@ -21,15 +21,15 @@
 from chemprop.parsing import add_predict_args, add_train_args, modify_predict_args, modify_train_args
 from chemprop.train.make_predictions import make_predictions
 from chemprop.train.run_training import run_training
-from chemprop.utils import create_logger, load_task_names
+from chemprop.utils import create_logger, load_task_names, makedirs
 
 TEMP_FOLDER = TemporaryDirectory()
 
 app = Flask(__name__)
 app.config['DATA_FOLDER'] = 'web_data'
-os.makedirs(app.config['DATA_FOLDER'], exist_ok=True)
+makedirs(app.config['DATA_FOLDER'])
 app.config['CHECKPOINT_FOLDER'] = 'web_checkpoints'
-os.makedirs(app.config['CHECKPOINT_FOLDER'], exist_ok=True)
+makedirs(app.config['CHECKPOINT_FOLDER'])
 app.config['TEMP_FOLDER'] = TEMP_FOLDER.name
 app.config['SMILES_FILENAME'] = 'smiles.csv'
 app.config['PREDICTIONS_FILENAME'] = 'predictions.csv'