Fix for systems without NVSMI

This commit is contained in:
iperov 2019-01-13 16:38:17 +04:00
parent 4625bcec1c
commit 12383570e8
3 changed files with 60 additions and 58 deletions

View file

@ -64,10 +64,12 @@ class ExtractSubprocessor(SubprocessorBase):
if (type == 'rects' or type == 'landmarks'): if (type == 'rects' or type == 'landmarks'):
if multi_gpu: if multi_gpu:
devices = nnlib.device.getDevicesWithAtLeastTotalMemoryGB(2) devices = nnlib.device.getDevicesWithAtLeastTotalMemoryGB(2)
else:
if not multi_gpu or len(devices) == 0:
devices = [nnlib.device.getBestDeviceIdx()] devices = [nnlib.device.getBestDeviceIdx()]
if devices[0] == -1:
devices = [] if len(devices) == 0 or devices[0] == -1:
devices = [0]
devices = [ (idx, nnlib.device.getDeviceName(idx), nnlib.device.getDeviceVRAMTotalGb(idx) ) for idx in devices] devices = [ (idx, nnlib.device.getDeviceName(idx), nnlib.device.getDeviceVRAMTotalGb(idx) ) for idx in devices]
@ -86,34 +88,30 @@ class ExtractSubprocessor(SubprocessorBase):
'detector': self.detector} 'detector': self.detector}
if not self.cpu_only: if not self.cpu_only:
devices = self.get_devices_for_type(self.type, self.multi_gpu) for (device_idx, device_name, device_total_vram_gb) in self.get_devices_for_type(self.type, self.multi_gpu):
if len(devices) != 0: num_processes = 1
for (device_idx, device_name, device_total_vram_gb) in devices: if not self.manual and self.type == 'rects' and self.detector == 'mt':
num_processes = 1 num_processes = int ( max (1, device_total_vram_gb / 2) )
if not self.manual and self.type == 'rects' and self.detector == 'mt':
num_processes = int ( max (1, device_total_vram_gb / 2) ) for i in range(0, num_processes ):
client_dict = base_dict.copy()
for i in range(0, num_processes ): client_dict['device_idx'] = device_idx
client_dict = base_dict.copy() client_dict['device_name'] = device_name if num_processes == 1 else '%s #%d' % (device_name,i)
client_dict['device_idx'] = device_idx client_dict['device_type'] = 'GPU'
client_dict['device_name'] = device_name if num_processes == 1 else '%s #%d' % (device_name,i)
client_dict['device_type'] = 'GPU' yield client_dict['device_name'], {}, client_dict
else:
yield client_dict['device_name'], {}, client_dict num_processes = 1
return if not self.manual and self.type == 'rects' and self.detector == 'mt':
print ("No capable GPU's found, falling back to CPU mode.") num_processes = int ( max (1, multiprocessing.cpu_count() / 2 ) )
num_processes = 1
if not self.manual and self.type == 'rects' and self.detector == 'mt':
num_processes = int ( max (1, multiprocessing.cpu_count() / 2 ) )
for i in range(0, num_processes ):
client_dict = base_dict.copy()
client_dict['device_idx'] = 0
client_dict['device_name'] = 'CPU' if num_processes == 1 else 'CPU #%d' % (i),
client_dict['device_type'] = 'CPU'
yield client_dict['device_name'], {}, client_dict for i in range(0, num_processes ):
client_dict = base_dict.copy()
client_dict['device_idx'] = 0
client_dict['device_name'] = 'CPU' if num_processes == 1 else 'CPU #%d' % (i),
client_dict['device_type'] = 'CPU'
yield client_dict['device_name'], {}, client_dict
#override #override
def get_no_process_started_message(self): def get_no_process_started_message(self):

View file

@ -24,7 +24,7 @@ class devicelib:
**in_options): **in_options):
self.use_fp16 = use_fp16 self.use_fp16 = use_fp16
if cpu_only or not devicelib.hasNVML(): if cpu_only:
self.cpu_only = True self.cpu_only = True
else: else:
self.force_best_gpu_idx = force_best_gpu_idx self.force_best_gpu_idx = force_best_gpu_idx
@ -34,30 +34,37 @@ class devicelib:
self.allow_growth = allow_growth self.allow_growth = allow_growth
self.gpu_idxs = [] self.gpu_idxs = []
if force_gpu_idxs is not None:
for idx in force_gpu_idxs.split(','): if not devicelib.hasNVML():
idx = int(idx) self.gpu_idxs = [0]
if devicelib.isValidDeviceIdx(idx): self.gpu_total_vram_gb = 2
self.gpu_idxs.append(idx) self.gpu_names += ['Generic GeForce GPU']
self.gpu_compute_caps += [ 50 ]
else: else:
gpu_idx = force_best_gpu_idx if (force_best_gpu_idx >= 0 and devicelib.isValidDeviceIdx(force_best_gpu_idx)) else devicelib.getBestDeviceIdx() if not choose_worst_gpu else devicelib.getWorstDeviceIdx() if force_gpu_idxs is not None:
if gpu_idx != -1: for idx in force_gpu_idxs.split(','):
if self.multi_gpu: idx = int(idx)
self.gpu_idxs = devicelib.getDeviceIdxsEqualModel( gpu_idx ) if devicelib.isValidDeviceIdx(idx):
if len(self.gpu_idxs) <= 1: self.gpu_idxs.append(idx)
self.multi_gpu = False else:
else: gpu_idx = force_best_gpu_idx if (force_best_gpu_idx >= 0 and devicelib.isValidDeviceIdx(force_best_gpu_idx)) else devicelib.getBestDeviceIdx() if not choose_worst_gpu else devicelib.getWorstDeviceIdx()
self.gpu_idxs = [gpu_idx] if gpu_idx != -1:
if self.multi_gpu:
self.cpu_only = (len(self.gpu_idxs) == 0) self.gpu_idxs = devicelib.getDeviceIdxsEqualModel( gpu_idx )
if len(self.gpu_idxs) <= 1:
if not self.cpu_only: self.multi_gpu = False
self.gpu_total_vram_gb = devicelib.getDeviceVRAMTotalGb ( self.gpu_idxs[0] ) else:
self.gpu_names = [] self.gpu_idxs = [gpu_idx]
self.gpu_compute_caps = []
for gpu_idx in self.gpu_idxs: self.cpu_only = (len(self.gpu_idxs) == 0)
self.gpu_names += [devicelib.getDeviceName(gpu_idx)]
self.gpu_compute_caps += [ devicelib.getDeviceComputeCapability ( gpu_idx ) ] if not self.cpu_only:
self.gpu_total_vram_gb = devicelib.getDeviceVRAMTotalGb ( self.gpu_idxs[0] )
self.gpu_names = []
self.gpu_compute_caps = []
for gpu_idx in self.gpu_idxs:
self.gpu_names += [devicelib.getDeviceName(gpu_idx)]
self.gpu_compute_caps += [ devicelib.getDeviceComputeCapability ( gpu_idx ) ]
@staticmethod @staticmethod
def hasNVML(): def hasNVML():
@ -215,7 +222,7 @@ class devicelib:
@staticmethod @staticmethod
def getDeviceName (idx): def getDeviceName (idx):
result = '' result = 'Generic GeForce GPU'
try: try:
nvmlInit() nvmlInit()
if idx < nvmlDeviceGetCount(): if idx < nvmlDeviceGetCount():

View file

@ -129,9 +129,6 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
if nnlib.tf is not None: if nnlib.tf is not None:
return nnlib.code_import_tf return nnlib.code_import_tf
if not nnlib.device.hasNVML():
print ("nvml.dll not found. Reinstall Geforce video drivers.")
if 'TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1': if 'TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1':
suppressor = std_utils.suppress_stdout_stderr().__enter__() suppressor = std_utils.suppress_stdout_stderr().__enter__()
else: else: