updates
This commit is contained in:
parent
1d49e66580
commit
3825e99ee3
|
@ -8,19 +8,21 @@ def init_seeds(seed=0):
|
||||||
|
|
||||||
|
|
||||||
def select_device(force_cpu=False):
|
def select_device(force_cpu=False):
|
||||||
if force_cpu:
|
cuda = False if force_cpu else torch.cuda.is_available()
|
||||||
cuda = False
|
device = torch.device('cuda:0' if cuda else 'cpu')
|
||||||
device = torch.device('cpu')
|
|
||||||
else:
|
|
||||||
cuda = torch.cuda.is_available()
|
|
||||||
device = torch.device('cuda:0' if cuda else 'cpu')
|
|
||||||
|
|
||||||
if torch.cuda.device_count() > 1:
|
if not cuda:
|
||||||
device = torch.device('cuda' if cuda else 'cpu')
|
print('Using CPU')
|
||||||
print('Found %g GPUs' % torch.cuda.device_count())
|
if cuda:
|
||||||
# print('Multi-GPU Issue: https://github.com/ultralytics/yolov3/issues/21')
|
c = 1024 ** 2 # bytes to MB
|
||||||
# torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available
|
ng = torch.cuda.device_count()
|
||||||
# print('Using ', torch.cuda.device_count(), ' GPUs')
|
x = [torch.cuda.get_device_properties(i) for i in range(ng)]
|
||||||
|
print("Using CUDA device0 _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
|
||||||
|
(x[0].name, x[0].total_memory / c))
|
||||||
|
if ng > 0:
|
||||||
|
# torch.cuda.set_device(0) # OPTIONAL: Set GPU ID
|
||||||
|
for i in range(1, ng):
|
||||||
|
print(" device%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
|
||||||
|
(i, x[i].name, x[i].total_memory / c))
|
||||||
|
|
||||||
print('Using %s %s\n' % (device.type, torch.cuda.get_device_properties(0) if cuda else ''))
|
|
||||||
return device
|
return device
|
||||||
|
|
|
@ -11,7 +11,7 @@ import torch.nn as nn
|
||||||
|
|
||||||
from utils import torch_utils
|
from utils import torch_utils
|
||||||
|
|
||||||
matplotlib.rc('font', **{'family': 'normal', 'size': 11})
|
matplotlib.rc('font', **{'family': 'normal', 'size': 12})
|
||||||
|
|
||||||
# Set printoptions
|
# Set printoptions
|
||||||
torch.set_printoptions(linewidth=1320, precision=5, profile='long')
|
torch.set_printoptions(linewidth=1320, precision=5, profile='long')
|
||||||
|
|
Loading…
Reference in New Issue