Merge pull request #49 from ecrl/dev

Additions to training runtime

ecnet/__init__.py

@@ -1 +1,2 @@
 from .model import ECNet
+__version__ = '4.1.0'

ecnet/model.py

@@ -61,7 +61,8 @@ def fit(self, smiles: List[str] = None, target_vals: List[List[float]] = None,
             dataset: QSPRDataset = None, backend: str = 'padel', batch_size: int = 32,
             epochs: int = 100, lr_decay: float = 0.0, valid_size: float = 0.0,
             valid_eval_iter: int = 1, patience: int = 16, verbose: int = 0,
-            random_state: int = None, **kwargs) -> Tuple[List[float], List[float]]:
+            random_state: int = None, shuffle: bool = False,
+            **kwargs) -> Tuple[List[float], List[float]]:
         """
         fit: fits ECNet to either (1) SMILES and target values, or (2) a pre-loaded QSPRDataset;
         the training process utilizes the Adam optimization algorithm, MSE loss, ReLU activation
@@ -90,6 +91,8 @@ def fit(self, smiles: List[str] = None, target_vals: List[List[float]] = None,
             verbose (int, optional): if > 0, will print every `this` epochs; default = 0
             random_state (int, optional): random_state used by sklearn.model_selection.
                 train_test_split; default = None
+            shuffle (bool, optional): if True, shuffles training/validation data between epochs;
+                default = False; random_state should be None
             **kwargs: arguments accepted by torch.optim.Adam (i.e. learning rate, beta values)
 
         Returns:
@@ -136,6 +139,18 @@ def fit(self, smiles: List[str] = None, target_vals: List[List[float]] = None,
             if not CBO.on_epoch_begin(epoch):
                 break
 
+            if shuffle:
+                index_train, index_valid = train_test_split(
+                    [i for i in range(len(dataset))], test_size=valid_size,
+                    random_state=random_state
+                )
+                dataloader_train = DataLoader(
+                    Subset(dataset, index_train), batch_size=batch_size, shuffle=True
+                )
+                dataloader_valid = DataLoader(
+                    Subset(dataset, index_valid), batch_size=len(index_valid), shuffle=True
+                )
+
             train_loss = 0.0
             self.train()
 

ecnet/tasks/__init__.py

@@ -1,2 +1,3 @@
 from .feature_selection import select_rfr
-from .parameter_tuning import tune_batch_size, tune_model_architecture, tune_training_parameters
+from .parameter_tuning import N_TESTS, CONFIG, tune_batch_size, tune_model_architecture,\
+    tune_training_parameters

ecnet/tasks/parameter_tuning.py

@@ -1,17 +1,19 @@
 from ecabc import ABC
 from sklearn.metrics import median_absolute_error
 from copy import deepcopy
+import numpy as np
 
 from ..model import ECNet
 from ..datasets.structs import QSPRDataset
 
 from typing import Iterable
 
+N_TESTS = 10
 
 CONFIG = {
     'training_params_range': {
-        'lr': (0.0, 0.05),
-        'lr_decay': (0.0, 0.0001)
+        'lr': (1e-16, 0.05),
+        'lr_decay': (1e-16, 0.0001)
     },
     'architecture_params_range': {
         'hidden_dim': (1, 1024),
@@ -38,6 +40,8 @@ def _get_kwargs(**kwargs):
         'eval_ds': kwargs.get('eval_ds'),
         'epochs': kwargs.get('epochs', 100),
         'batch_size': kwargs.get('batch_size', 32),
+        'valid_size': kwargs.get('valid_size', 0.2),
+        'patience': kwargs.get('patience', 32),
         'lr_decay': kwargs.get('lr_decay', 0.0),
         'lr': kwargs.get('lr', 0.001),
         'beta_1': kwargs.get('beta_1', 0.9),
@@ -54,6 +58,8 @@ def _get_kwargs(**kwargs):
 def _evaluate_model(trial_spec: dict) -> float:
     """
     Training sub-function for cost functions _cost_batch_size, _cost_arch, _cost_train_hp;
+    Each model configuration is tested ecnet.tasks.parameter_tuning.N_TESTS times, average
+    median absolute error across all tests returned; default 10 tests per configuration
 
     Args:
         trial_spec (dict): all relevant parameters for this training trial
@@ -62,25 +68,32 @@ def _evaluate_model(trial_spec: dict) -> float:
         float: median absolute error for dataset being evaluated (trial_spec['eval_ds'])
     """
 
-    model = deepcopy(trial_spec['model'])
-    model._hidden_dim = trial_spec['hidden_dim']
-    model._n_hidden = trial_spec['n_hidden']
-    model._dropout = trial_spec['dropout']
-    model._construct()
-    model.fit(
-        dataset=trial_spec['train_ds'],
-        epochs=trial_spec['epochs'],
-        batch_size=trial_spec['batch_size'],
-        lr_decay=trial_spec['lr_decay'],
-        lr=trial_spec['lr_decay'],
-        betas=(trial_spec['beta_1'], trial_spec['beta_2']),
-        eps=trial_spec['eps'],
-        weight_decay=trial_spec['weight_decay'],
-        amsgrad=trial_spec['amsgrad']
+    model = ECNet(
+        trial_spec['train_ds'].desc_vals.shape[1],
+        trial_spec['train_ds'].target_vals.shape[1],
+        trial_spec['hidden_dim'],
+        trial_spec['n_hidden'],
+        trial_spec['dropout']
     )
-    yhat_eval = model(trial_spec['eval_ds'].desc_vals).detach().numpy()
-    y_eval = trial_spec['eval_ds'].target_vals
-    return median_absolute_error(y_eval, yhat_eval)
+    maes = []
+    for _ in range(N_TESTS):
+        model._construct()
+        model.fit(
+            dataset=trial_spec['train_ds'],
+            epochs=trial_spec['epochs'],
+            batch_size=trial_spec['batch_size'],
+            patience=trial_spec['patience'],
+            lr_decay=trial_spec['lr_decay'],
+            lr=trial_spec['lr_decay'],
+            betas=(trial_spec['beta_1'], trial_spec['beta_2']),
+            eps=trial_spec['eps'],
+            weight_decay=trial_spec['weight_decay'],
+            amsgrad=trial_spec['amsgrad']
+        )
+        yhat_eval = model(trial_spec['eval_ds'].desc_vals).detach().numpy()
+        y_eval = trial_spec['eval_ds'].target_vals
+        maes.append(median_absolute_error(y_eval, yhat_eval))
+    return np.mean(maes)
 
 
 def _cost_batch_size(vals: Iterable[float], **kwargs) -> float:
@@ -100,20 +113,43 @@ def _cost_batch_size(vals: Iterable[float], **kwargs) -> float:
     return _evaluate_model(trial_spec)
 
 
-def tune_batch_size(n_bees: int, n_iter: int, n_processes: int = 1, **kwargs) -> dict:
+def tune_batch_size(n_bees: int, n_iter: int, dataset_train: QSPRDataset,
+                    dataset_eval: QSPRDataset, n_processes: int = 1,
+                    **kwargs) -> dict:
     """
-    Tunes the batch size during training
+    Tunes the batch size during training; additional **kwargs can include any in:
+        [
+            # ECNet parameters
+            'epochs' (default 100),
+            'valid_size' (default 0.2),
+            'patience' (default 32),
+            'lr_decay' (default 0.0),
+            'hidden_dim' (default 128),
+            'n_hidden' (default 2),
+            'dropout': (default 0.0),
+            # Adam optim. alg. arguments
+            'lr' (default 0.001),
+            'beta_1' (default 0.9),
+            'beta_2' (default 0.999),
+            'eps' (default 1e-8),
+            'weight_decay' (default 0.0),
+            'amsgrad' (default False)
+        ]
 
     Args:
         n_bees (int): number of employer bees to use in ABC algorithm
         n_iter (int): number of iterations, or "search cycles", for ABC algorithm
-        n_processes (int): if > 1, uses multiprocessing when evaluating at an iteration
-        **kwargs: arguments passed to _cost_batch_size
+        dataset_train (QSPRDataset): dataset used to train evaluation models
+        dataset_eval (QSPRDataset): dataset used for evaluation
+        n_processes (int, optional): if > 1, uses multiprocessing when evaluating at an iteration
+        **kwargs: additional arguments 
 
     Returns:
-        dict: {'batch_size': tuned batch size}
+        dict: {'batch_size': int}
     """
 
+    kwargs['train_ds'] = dataset_train
+    kwargs['eval_ds'] = dataset_eval
     abc = ABC(n_bees, _cost_batch_size, num_processes=n_processes, obj_fn_args=kwargs)
     abc.add_param(1, len(kwargs.get('train_ds').desc_vals), name='batch_size')
     abc.initialize()
@@ -144,20 +180,42 @@ def _cost_arch(vals, **kwargs):
     return _evaluate_model(trial_spec)
 
 
-def tune_model_architecture(n_bees: int, n_iter: int, n_processes: int = 1, **kwargs) -> dict:
+def tune_model_architecture(n_bees: int, n_iter: int, dataset_train: QSPRDataset,
+                            dataset_eval: QSPRDataset, n_processes: int = 1,
+                            **kwargs) -> dict:
     """
-    Tunes the NN's architecture
+    Tunes model architecture parameters (number of hidden layers, neurons per hidden layer, neuron
+    dropout); additional **kwargs can include any in:
+        [
+            # ECNet parameters
+            'epochs' (default 100),
+            'batch_size' (default 32),
+            'valid_size' (default 0.2),
+            'patience' (default 32),
+            'lr_decay' (default 0.0),
+            # Adam optim. alg. arguments
+            'lr' (default 0.001),
+            'beta_1' (default 0.9),
+            'beta_2' (default 0.999),
+            'eps' (default 1e-8),
+            'weight_decay' (default 0.0),
+            'amsgrad' (default False)
+        ]
 
     Args:
         n_bees (int): number of employer bees to use in ABC algorithm
         n_iter (int): number of iterations, or "search cycles", for ABC algorithm
-        n_processes (int): if > 1, uses multiprocessing when evaluating at an iteration
-        **kwargs: arguments passed to _cost_batch_size
+        dataset_train (QSPRDataset): dataset used to train evaluation models
+        dataset_eval (QSPRDataset): dataset used for evaluation
+        n_processes (int, optional): if > 1, uses multiprocessing when evaluating at an iteration
+        **kwargs: additional arguments 
 
     Returns:
-        dict: {'batch_size': opt_val, 'n_hidden': opt_val, 'dropout': opt_val}
+        dict: {'hidden_dim': int, 'n_hidden': int, 'dropout': float}
     """
 
+    kwargs['train_ds'] = dataset_train
+    kwargs['eval_ds'] = dataset_eval
     abc = ABC(n_bees, _cost_arch, num_processes=n_processes, obj_fn_args=kwargs)
     abc.add_param(CONFIG['architecture_params_range']['hidden_dim'][0],
                   CONFIG['architecture_params_range']['hidden_dim'][1], name='hidden_dim')
@@ -195,20 +253,42 @@ def _cost_train_hp(vals, **kwargs):
     return _evaluate_model(trial_spec)
 
 
-def tune_training_parameters(n_bees: int, n_iter: int, n_processes: int = 1, **kwargs) -> dict:
+def tune_training_parameters(n_bees: int, n_iter: int, dataset_train: QSPRDataset,
+                             dataset_eval: QSPRDataset, n_processes: int = 1,
+                             **kwargs) -> dict:
     """
-    Tunes the NN's training parameters (Adam optim. fn.)
+    Tunes learning rate, learning rate decay; additional **kwargs can include any in:
+        [
+            # ECNet parameters
+            'epochs' (default 100),
+            'batch_size' (default 32),
+            'valid_size' (default 0.2),
+            'patience' (default 32),
+            'hidden_dim' (default 128),
+            'n_hidden' (default 2),
+            'dropout': (default 0.0),
+            # Adam optim. alg. arguments
+            'beta_1' (default 0.9),
+            'beta_2' (default 0.999),
+            'eps' (default 1e-8),
+            'weight_decay' (default 0.0),
+            'amsgrad' (default False)
+        ]
 
     Args:
         n_bees (int): number of employer bees to use in ABC algorithm
         n_iter (int): number of iterations, or "search cycles", for ABC algorithm
-        n_processes (int): if > 1, uses multiprocessing when evaluating at an iteration
-        **kwargs: arguments passed to _cost_batch_size
+        dataset_train (QSPRDataset): dataset used to train evaluation models
+        dataset_eval (QSPRDataset): dataset used for evaluation
+        n_processes (int, optional): if > 1, uses multiprocessing when evaluating at an iteration
+        **kwargs: additional arguments 
 
     Returns:
-        dict: {'lr': opt_val, 'lr_decay': opt_val}
+        dict: {'lr': float, 'lr_decay': float}
     """
 
+    kwargs['train_ds'] = dataset_train
+    kwargs['eval_ds'] = dataset_eval
     abc = ABC(n_bees, _cost_train_hp, num_processes=n_processes, obj_fn_args=kwargs)
     abc.add_param(CONFIG['training_params_range']['lr'][0],
                   CONFIG['training_params_range']['lr'][1], name='lr')

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='ecnet',
-    version='4.0.0',
+    version='4.1.0',
     description='Fuel property prediction using QSPR descriptors',
     url='https://github.com/ecrl/ecnet',
     author='Travis Kessler',

tests/test_all.py

@@ -257,7 +257,7 @@ def test_tune_batch_size(self):
         targets = [[5.0]]
         ds_eval = QSPRDataset(smiles, targets, backend=_BACKEND)
         model = ECNet(_N_DESC, 1, 5, 1)
-        res = tune_batch_size(1, 1, _N_PROCESSES, model=model, train_ds=ds_train, eval_ds=ds_eval)
+        res = tune_batch_size(1, 1, ds_train, ds_eval, _N_PROCESSES)
         self.assertTrue(1 <= res['batch_size'] <= len(ds_train.target_vals))
 
     def test_tune_model_architecture(self):
@@ -270,8 +270,7 @@ def test_tune_model_architecture(self):
         targets = [[5.0]]
         ds_eval = QSPRDataset(smiles, targets, backend=_BACKEND)
         model = ECNet(_N_DESC, 1, 5, 1)
-        res = tune_model_architecture(1, 1, _N_PROCESSES, model=model, train_ds=ds_train,
-                                      eval_ds=ds_eval)
+        res = tune_model_architecture(1, 1, ds_train, ds_eval, _N_PROCESSES)
         for k in list(res.keys()):
             self.assertTrue(res[k] >= CONFIG['architecture_params_range'][k][0])
             self.assertTrue(res[k] <= CONFIG['architecture_params_range'][k][1])
@@ -285,12 +284,8 @@ def test_tune_training_hps(self):
         smiles = ['CCCCC']
         targets = [[5.0]]
         ds_eval = QSPRDataset(smiles, targets, backend=_BACKEND)
-        model = ECNet(_N_DESC, 1, 5, 1)
-        res = tune_training_parameters(1, 1, _N_PROCESSES, model=model, train_ds=ds_train,
-                                       eval_ds=ds_eval)
+        res = tune_training_parameters(1, 1, ds_train, ds_eval, _N_PROCESSES)
         for k in list(res.keys()):
-            if k == 'betas':
-                continue
             self.assertTrue(res[k] >= CONFIG['training_params_range'][k][0])
             self.assertTrue(res[k] <= CONFIG['training_params_range'][k][1])