diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index ea2544f..01cebe0 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -64,10 +64,12 @@ class ExtractSubprocessor(SubprocessorBase): if (type == 'rects' or type == 'landmarks'): if multi_gpu: devices = nnlib.device.getDevicesWithAtLeastTotalMemoryGB(2) - else: + + if not multi_gpu or len(devices) == 0: 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] @@ -86,34 +88,30 @@ class ExtractSubprocessor(SubprocessorBase): 'detector': self.detector} if not self.cpu_only: - devices = self.get_devices_for_type(self.type, self.multi_gpu) - if len(devices) != 0: - for (device_idx, device_name, device_total_vram_gb) in devices: - num_processes = 1 - 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() - client_dict['device_idx'] = device_idx - 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 - return - print ("No capable GPU's found, falling back to CPU mode.") - - 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' + for (device_idx, device_name, device_total_vram_gb) in self.get_devices_for_type(self.type, self.multi_gpu): + num_processes = 1 + 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() + client_dict['device_idx'] = device_idx + 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: + num_processes = 1 + if not self.manual and self.type == 'rects' and self.detector == 'mt': + num_processes = int ( max (1, multiprocessing.cpu_count() / 2 ) ) - 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 def get_no_process_started_message(self): diff --git a/nnlib/devicelib.py b/nnlib/devicelib.py index cd3c549..74eb3d3 100644 --- a/nnlib/devicelib.py +++ b/nnlib/devicelib.py @@ -24,7 +24,7 @@ class devicelib: **in_options): self.use_fp16 = use_fp16 - if cpu_only or not devicelib.hasNVML(): + if cpu_only: self.cpu_only = True else: self.force_best_gpu_idx = force_best_gpu_idx @@ -34,30 +34,37 @@ class devicelib: self.allow_growth = allow_growth self.gpu_idxs = [] - if force_gpu_idxs is not None: - for idx in force_gpu_idxs.split(','): - idx = int(idx) - if devicelib.isValidDeviceIdx(idx): - self.gpu_idxs.append(idx) + + if not devicelib.hasNVML(): + self.gpu_idxs = [0] + self.gpu_total_vram_gb = 2 + self.gpu_names += ['Generic GeForce GPU'] + self.gpu_compute_caps += [ 50 ] 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 gpu_idx != -1: - if self.multi_gpu: - self.gpu_idxs = devicelib.getDeviceIdxsEqualModel( gpu_idx ) - if len(self.gpu_idxs) <= 1: - self.multi_gpu = False - else: - self.gpu_idxs = [gpu_idx] - - self.cpu_only = (len(self.gpu_idxs) == 0) - - 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 ) ] + if force_gpu_idxs is not None: + for idx in force_gpu_idxs.split(','): + idx = int(idx) + if devicelib.isValidDeviceIdx(idx): + self.gpu_idxs.append(idx) + 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 gpu_idx != -1: + if self.multi_gpu: + self.gpu_idxs = devicelib.getDeviceIdxsEqualModel( gpu_idx ) + if len(self.gpu_idxs) <= 1: + self.multi_gpu = False + else: + self.gpu_idxs = [gpu_idx] + + self.cpu_only = (len(self.gpu_idxs) == 0) + + 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 def hasNVML(): @@ -215,7 +222,7 @@ class devicelib: @staticmethod def getDeviceName (idx): - result = '' + result = 'Generic GeForce GPU' try: nvmlInit() if idx < nvmlDeviceGetCount(): diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index 40c089a..3ce3fd8 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -129,9 +129,6 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator if nnlib.tf is not None: 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': suppressor = std_utils.suppress_stdout_stderr().__enter__() else: