From 65ed0d11222411e69ea02aaf89e42fa994df9f88 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 4 Mar 2019 15:46:25 +0100 Subject: [PATCH] updates --- train.py | 4 ++-- utils/torch_utils.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/train.py b/train.py index 86253e54..59ac48bc 100644 --- a/train.py +++ b/train.py @@ -75,8 +75,8 @@ def train( load_darknet_weights(model, weights + 'yolov3-tiny.conv.15') cutoff = 15 - # if torch.cuda.device_count() > 1: - # model = nn.DataParallel(model) + if torch.cuda.device_count() > 1: + model = nn.DataParallel(model) model.to(device).train() # Set optimizer diff --git a/utils/torch_utils.py b/utils/torch_utils.py index c1629f3a..dc7a4e69 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -15,8 +15,9 @@ def select_device(force_cpu=False): cuda = torch.cuda.is_available() device = torch.device('cuda:0' if cuda else 'cpu') - # if torch.cuda.device_count() > 1: - # print('WARNING Using GPU0 Only: https://github.com/ultralytics/yolov3/issues/21') + if torch.cuda.device_count() > 1: + print('Found %g GPUs' % torch.cuda.device_count()) + print('WARNING Using GPU0 Only: https://github.com/ultralytics/yolov3/issues/21') # torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available # # print('Using ', torch.cuda.device_count(), ' GPUs')