Specify GPU inds to use for multi-gpu testing

Reviewed By: rbgirshick

Differential Revision: D6784702

fbshipit-source-id: b5f81d858f9e9d33c048c39f849d0e29e8db3cb1

lib/utils/subprocess.py

@@ -49,10 +49,19 @@ def process_in_parallel(tag, total_range_size, binary, output_dir):
     subprocess_env = os.environ.copy()
     processes = []
     subinds = np.array_split(range(total_range_size), cfg.NUM_GPUS)
-    for i in range(cfg.NUM_GPUS):
+    # Determine GPUs to use
+    cuda_visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES')
+    if cuda_visible_devices:
+        gpu_inds = map(int, cuda_visible_devices.split(','))
+        assert -1 not in gpu_inds, \
+            'Hiding GPU indices using the \'-1\' index is not supported'
+    else:
+        gpu_inds = range(cfg.NUM_GPUS)
+    # Run the binary in cfg.NUM_GPUS subprocesses
+    for i, gpu_ind in enumerate(gpu_inds):
         start = subinds[i][0]
         end = subinds[i][-1] + 1
-        subprocess_env['CUDA_VISIBLE_DEVICES'] = str(i)
+        subprocess_env['CUDA_VISIBLE_DEVICES'] = str(gpu_ind)
         cmd = '{binary} --range {start} {end} --cfg {cfg_file} NUM_GPUS 1'
         cmd = cmd.format(
             binary=shlex_quote(binary),