Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-03-04 15:46:25 +01:00
parent 66211d1147
commit 65ed0d1122
2 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
train.py
View File

@ -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

5
utils/torch_utils.py
View File

@ -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')