Merge pull request #60 from francois-drielsma/develop

Add option to specify visible devices manually

Author: francois-drielsma

bin/run.py

@@ -107,22 +107,17 @@ def main(config, source, source_list, output, n, nskip, detect_anomaly,
         if not 'train' in cfg['base']:
             raise KeyError("--weight_prefix flag provided: must specify "
                            "`train` in the `base` block.")
-        cfg['base']['train']['weight_prefix']=weight_prefix
+        cfg['base']['train']['weight_prefix'] = weight_prefix
 
     if weight_path is not None:
-        cfg['model']['weight_path']=weight_path
+        cfg['model']['weight_path'] = weight_path
 
     # Turn on PyTorch anomaly detection, if requested
     if detect_anomaly is not None:
         assert 'model' in cfg, (
                 "There is no model to detect anomalies for, add `model` block.")
         cfg['model']['detect_anomaly'] = detect_anomaly
 
-    # If the -1 option for GPUs is selected, expose the process to all GPUs
-    if os.environ.get('CUDA_VISIBLE_DEVICES') is not None \
-            and cfg['base'].get('gpus', '') == '-1':
-        cfg['base']['gpus'] = os.environ.get('CUDA_VISIBLE_DEVICES')
-
     # Execute train/validation process
     run(cfg)
 

spine/driver.py

@@ -153,8 +153,8 @@ def process_config(self, io, base=None, model=None, build=None,
         ana : dict, optional
             Analysis script configurationdictionary
         rank : int, optional
-            Rank of the GPU. If not specified, the model will be run on CPU if
-            `world_size` is 0 and GPU is `world_size` is > 0.
+            Rank of the GPU. The model will be run on CPU if `world_size` is not
+            specified or 0 and on GPU is `world_size` is > 0.
 
         Returns
         -------
@@ -170,9 +170,15 @@ def process_config(self, io, base=None, model=None, build=None,
         logger.setLevel(verbosity.upper())
 
         # Set GPUs visible to CUDA
-        world_size = base.get('world_size', 0)
-        os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
-            [str(i) for i in range(world_size)])
+        gpus = base.get('gpus', None)
+        if gpus is not None:
+            os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
+                    [str(i) for i in gpus])
+
+        elif not os.environ['CUDA_VISIBLE_DEVICES']:
+            world_size = base.get('world_size', 0)
+            os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
+                [str(i) for i in range(world_size)])
 
         # If the seed is not set for the sampler, randomize it. This is done
         # here to keep a record of the seeds provided to the samplers
@@ -224,7 +230,7 @@ def process_config(self, io, base=None, model=None, build=None,
         # Return updated configuration
         return base, io, model, build, post, ana
 
-    def initialize_base(self, seed, dtype='float32', world_size=0,
+    def initialize_base(self, seed, dtype='float32', world_size=None, gpus=None,
                         log_dir='logs', prefix_log=False, overwrite_log=False,
                         parent_path=None, iterations=None, epochs=None,
                         unwrap=False, rank=None, log_step=1, distributed=False,
@@ -237,8 +243,10 @@ def initialize_base(self, seed, dtype='float32', world_size=0,
             Random number generator seed
         dtype : str, default 'float32'
             Data type of the model parameters and input data
-        world_size : int, default 0
+        world_size : int, optional
             Number of GPUs to use in the underlying model
+        gpus : List[int], optional
+            List of indexes of GPUs to expose to the model
         log_dir : str, default 'logs'
             Path to the directory where the logs will be written to
         prefix_log : bool, default False
@@ -279,6 +287,16 @@ def initialize_base(self, seed, dtype='float32', world_size=0,
         numba_seed(seed)
         torch.manual_seed(seed)
 
+        # Check on the number of GPUs to use
+        if gpus is not None:
+            assert world_size is None or len(gpus) == world_size, (
+                    f"The number of visible GPUs ({len(gpus)}) is not "
+                    f"compatible with the world size ({world_size}).")
+            world_size = len(gpus)
+
+        elif world_size is None:
+            world_size = 0
+
         # Set up the device the model will run on
         if rank is None and world_size > 0:
             assert world_size < 2, (