From c9a9e5c5256aa6cafd0a6674e3ce676d35e6d3b4 Mon Sep 17 00:00:00 2001 From: James Betker Date: Thu, 1 Oct 2020 17:24:50 -0600 Subject: [PATCH] Prompt user for gpu_id if multiple gpus are detected --- codes/train.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/codes/train.py b/codes/train.py index 7c6f360e..eef9f237 100644 --- a/codes/train.py +++ b/codes/train.py @@ -32,7 +32,7 @@ def init_dist(backend='nccl', **kwargs): def main(): #### options parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_imgset_spsr7_multiloss.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_imgset_ssgr.yml') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args() @@ -61,6 +61,12 @@ def main(): opt['dist'] = False rank = -1 print('Disabled distributed training.') + if torch.cuda.device_count() > 1: + gpu = input('I noticed you have multiple GPUs. Starting two jobs on the same GPU sucks. Please confirm which GPU' + 'you want to use. Press enter to use the specified one [%i]' % (opt['gpu_ids'])) + if gpu: + opt['gpu_ids'] = [int(gpu)] + else: opt['dist'] = True init_dist()