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

updates

This commit is contained in:
Glenn Jocher 2019-02-25 13:47:51 +01:00
parent 8af70386e8
commit d2cd49f059
1 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
utils/torch_utils.py
View File

@ -3,21 +3,22 @@ import torch
def init_seeds(seed=0): def init_seeds(seed=0):
torch.manual_seed(seed) torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed) torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) torch.cuda.manual_seed_all(seed)
def select_device(force_cpu=False): def select_device(force_cpu=False):
if force_cpu: if force_cpu:
cuda = False
device = torch.device('cpu') device = torch.device('cpu')
else: else:
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')
if torch.cuda.device_count() > 1: if torch.cuda.device_count() > 1:
print('WARNING Using GPU0 Only. Multi-GPU issue: https://github.com/ultralytics/yolov3/issues/21') print('WARNING Using GPU0 Only: https://github.com/ultralytics/yolov3/issues/21')
torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available
# print('Using ', torch.cuda.device_count(), ' GPUs') # print('Using ', torch.cuda.device_count(), ' GPUs')
print('Using ' + str(device) + '\n') print('Using %s\n%s' % (device.type, torch.cuda.get_device_properties(0) if cuda else ''))
return device return device