github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-12 08:07:03 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactoring

Browse source
This commit is contained in:
iperov 2018-12-22 19:37:25 +04:00
parent 0d7387165a
commit 44798c2b85
1 changed files with 34 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

50
mainscripts/Extractor.py
View file

@ -60,10 +60,8 @@ class ExtractSubprocessor(SubprocessorBase):
cv2.setMouseCallback(self.wnd_name, onMouse, self.param) cv2.setMouseCallback(self.wnd_name, onMouse, self.param)
def get_devices_for_type (self, type, multi_gpu, cpu_only): def get_devices_for_type (self, type, multi_gpu):
if cpu_only: if (type == 'rects' or type == 'landmarks'):
devices = [ (0, 'CPU', 0 ) ]
elif (type == 'rects' or type == 'landmarks'):
if not multi_gpu: if not multi_gpu:
devices = [gpufmkmgr.getBestDeviceIdx()] devices = [gpufmkmgr.getBestDeviceIdx()]
else: else:
@ -77,25 +75,39 @@ class ExtractSubprocessor(SubprocessorBase):
#override #override
def process_info_generator(self): def process_info_generator(self):
for (device_idx, device_name, device_total_vram_gb) in self.get_devices_for_type(self.type, self.multi_gpu, self.cpu_only): base_dict = {'type' : self.type,
'image_size': self.image_size,
'face_type': self.face_type,
'debug': self.debug,
'output_dir': str(self.output_path),
'detector': self.detector}
if self.cpu_only:
num_processes = 1 num_processes = 1
if not self.manual and self.type == 'rects' and self.detector == 'mt': if not self.manual and self.type == 'rects' and self.detector == 'mt':
if self.cpu_only: num_processes = int ( max (1, multiprocessing.cpu_count() / 2 ) )
num_processes = int ( max (1, multiprocessing.cpu_count() / 2 ) )
else:
num_processes = int ( max (1, device_total_vram_gb / 2) )
for i in range(0, num_processes ): for i in range(0, num_processes ):
device_name_for_process = device_name if num_processes == 1 else '%s #%d' % (device_name,i) client_dict = base_dict.copy()
yield device_name_for_process, {}, {'type' : self.type, client_dict['device_idx'] = 0
'device_idx' : device_idx, client_dict['device_name'] = 'CPU' if num_processes == 1 else 'CPU #%d' % (i),
'device_name' : device_name_for_process, client_dict['device_type'] = 'CPU'
'device_type' : 'CPU' if self.cpu_only else 'GPU',
'image_size': self.image_size, yield client_dict['device_name'], {}, client_dict
'face_type': self.face_type,
'debug': self.debug, else:
'output_dir': str(self.output_path), for (device_idx, device_name, device_total_vram_gb) in self.get_devices_for_type(self.type, self.multi_gpu):
'detector': self.detector} 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
#override #override
def get_no_process_started_message(self): def get_no_process_started_message(self):