set deterministic parameter in TrainModel (#181)

* set deterministic in TrainModel

* update library setup.py with PyTorch 1.9

* move logger init earlier

* fix output_dir init

Author: kathyxchen

README.md

@@ -12,7 +12,7 @@ We recommend using Selene with Python 3.6 or above.
 Package installation should only take a few minutes (less than 10 minutes, typically ~2-3 minutes) with any of these methods (conda, pip, source). 
 
 **First, install [PyTorch](https://pytorch.org/get-started/locally/).** If you have an NVIDIA GPU, install a version of PyTorch that supports it--Selene will run much faster with a discrete GPU. 
-The library is currently compatible with PyTorch versions between 0.4.1 and 1.4.0.
+The library is currently compatible with PyTorch versions between 0.4.1 and 1.9.
 We will continue to update Selene to be compatible with the latest version of PyTorch.
 
 ### Installing selene with [Anaconda](https://www.anaconda.com/download/) (for Linux):

selene_sdk/train_model.py

@@ -138,8 +138,14 @@ class TrainModel(object):
         using `torch.load`.
     use_scheduler : bool, optional
         Default is `True`. If `True`, learning rate scheduler is used to
-        reduce learning rate on plateau. PyTorch ReduceLROnPlateau scheduler 
+        reduce learning rate on plateau. PyTorch ReduceLROnPlateau scheduler
         with patience=16 and factor=0.8 is used.
+    deterministic : bool, optional
+        Default is `False`. If `True`, will set
+        `torch.backends.cudnn.deterministic` to True and
+        `torch.backends.cudnn.benchmark = False`. In Selene CLI,
+        if `random_seed` is set in the configuration YAML, Selene automatically
+        passes in `deterministic=True` to the TrainModel class.
 
     Attributes
     ----------
@@ -190,7 +196,8 @@ def __init__(self,
                  checkpoint_resume=None,
                  metrics=dict(roc_auc=roc_auc_score,
                               average_precision=average_precision_score),
-                 use_scheduler=True):
+                 use_scheduler=True,
+                 deterministic=False):
         """
         Constructs a new `TrainModel` object.
         """
@@ -212,6 +219,19 @@ def __init__(self,
 
         self._save_new_checkpoints = save_new_checkpoints_after_n_steps
 
+        os.makedirs(output_dir, exist_ok=True)
+        self.output_dir = output_dir
+
+        initialize_logger(
+            os.path.join(self.output_dir, "{0}.log".format(__name__)),
+            verbosity=logging_verbosity)
+
+        if deterministic:
+            logger.info("Setting deterministic = True for reproducibility.")
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+
+
         logger.info("Training parameters set: batch size {0}, "
                     "number of steps per 'epoch': {1}, "
                     "maximum number of steps: {2}".format(
@@ -233,13 +253,6 @@ def __init__(self,
             self.criterion.cuda()
             logger.debug("Set modules to use CUDA")
 
-        os.makedirs(output_dir, exist_ok=True)
-        self.output_dir = output_dir
-
-        initialize_logger(
-            os.path.join(self.output_dir, "{0}.log".format(__name__)),
-            verbosity=logging_verbosity)
-
         self._report_gt_feature_n_positives = report_gt_feature_n_positives
         self._metrics = metrics
         self._n_validation_samples = n_validation_samples

selene_sdk/utils/config_utils.py

@@ -181,6 +181,8 @@ def execute(operations, configs, output_dir):
                                   optimizer_kwargs=optim_kwargs)
             if output_dir is not None:
                 train_model_info.bind(output_dir=output_dir)
+            if "random_seed" in configs:
+                train_model_info.bind(deterministic=True)
 
             train_model = instantiate(train_model_info)
             # TODO: will find a better way to handle this in the future
@@ -341,8 +343,7 @@ def parse_configs_and_run(configs,
         np.random.seed(seed)
         torch.manual_seed(seed)
         torch.cuda.manual_seed_all(seed)
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
+        print("Setting random seed = {0}".format(seed))
     else:
         print("Warning: no random seed specified in config file. "
               "Using a random seed ensures results are reproducible.")

setup.py

@@ -63,7 +63,7 @@
         "scipy",
         "seaborn",
         "statsmodels",
-        "torch>=0.4.1, <=1.4.0",
+        "torch>=0.4.1, <=1.9",
     ],
     entry_points={
         'console_scripts': [