def onClientInitialize(self, client_dict):
self.safe_print('Running on %s.' % (client_dict['device_name']))
self.type = client_dict['type']
self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type']
self.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.output_path = Path(
client_dict['output_dir']) if 'output_dir' in client_dict.keys(
) else None
self.debug = client_dict['debug']
self.detector = client_dict['detector']
self.keras = None
self.tf = None
self.tf_session = None
self.e = None
if self.type == 'rects':
if self.detector is not None:
if self.detector == 'mt':
self.gpu_config = gpufmkmgr.GPUConfig(
cpu_only=self.cpu_only,
force_best_gpu_idx=self.device_idx,
allow_growth=True)
self.tf = gpufmkmgr.import_tf(self.gpu_config)
self.tf_session = gpufmkmgr.get_tf_session()
self.keras = gpufmkmgr.import_keras()
self.e = facelib.MTCExtractor(self.keras, self.tf,
self.tf_session)
elif self.detector == 'dlib':
self.dlib = gpufmkmgr.import_dlib(self.device_idx,
cpu_only=self.cpu_only)
self.e = facelib.DLIBExtractor(self.dlib)
self.e.__enter__()
elif self.type == 'landmarks':
self.gpu_config = gpufmkmgr.GPUConfig(
cpu_only=self.cpu_only,
force_best_gpu_idx=self.device_idx,
allow_growth=True)
self.tf = gpufmkmgr.import_tf(self.gpu_config)
self.tf_session = gpufmkmgr.get_tf_session()
self.keras = gpufmkmgr.import_keras()
self.e = facelib.LandmarksExtractor(self.keras)
self.e.__enter__()
elif self.type == 'final':
pass
return None
def onClientInitialize(self, client_dict):
print ('Running on %s.' % (client_dict['device_name']) )
self.device_idx = client_dict['device_idx']
self.device_name = client_dict['device_name']
self.converter = client_dict['converter']
self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None
self.alignments = client_dict['alignments']
self.debug = client_dict['debug']
import gpufmkmgr
#model process ate all GPU mem,
#so we cannot use GPU for any TF operations in converter processes (for example image_utils.TFLabConverter)
gpufmkmgr.set_prefer_GPUConfig ( gpufmkmgr.GPUConfig (cpu_only=True) )
return None
def __init__(self,
model_path,
training_data_src_path=None,
training_data_dst_path=None,
batch_size=0,
write_preview_history=False,
debug=False,
**in_options):
print("Loading model...")
self.model_path = model_path
self.model_data_path = Path(
self.get_strpath_storage_for_file('data.dat'))
self.training_data_src_path = training_data_src_path
self.training_data_dst_path = training_data_dst_path
self.src_images_paths = None
self.dst_images_paths = None
self.src_yaw_images_paths = None
self.dst_yaw_images_paths = None
self.src_data_generator = None
self.dst_data_generator = None
self.is_training_mode = (training_data_src_path is not None
and training_data_dst_path is not None)
self.batch_size = batch_size
self.write_preview_history = write_preview_history
self.debug = debug
self.supress_std_once = ('TF_SUPPRESS_STD' in os.environ.keys()
and os.environ['TF_SUPPRESS_STD'] == '1')
if self.model_data_path.exists():
model_data = pickle.loads(self.model_data_path.read_bytes())
self.epoch = model_data['epoch']
self.options = model_data['options']
self.loss_history = model_data[
'loss_history'] if 'loss_history' in model_data.keys() else []
self.sample_for_preview = model_data[
'sample_for_preview'] if 'sample_for_preview' in model_data.keys(
) else None
else:
self.epoch = 0
self.options = {}
self.loss_history = []
self.sample_for_preview = None
if self.write_preview_history:
self.preview_history_path = self.model_path / (
'%s_history' % (self.get_model_name()))
if not self.preview_history_path.exists():
self.preview_history_path.mkdir(exist_ok=True)
else:
if self.epoch == 0:
for filename in Path_utils.get_image_paths(
self.preview_history_path):
Path(filename).unlink()
self.gpu_config = gpufmkmgr.GPUConfig(allow_growth=False, **in_options)
self.gpu_total_vram_gb = self.gpu_config.gpu_total_vram_gb
if self.epoch == 0:
#first run
self.options['created_vram_gb'] = self.gpu_total_vram_gb
self.created_vram_gb = self.gpu_total_vram_gb
else:
#not first run
if 'created_vram_gb' in self.options.keys():
self.created_vram_gb = self.options['created_vram_gb']
else:
self.options['created_vram_gb'] = self.gpu_total_vram_gb
self.created_vram_gb = self.gpu_total_vram_gb
self.tf = gpufmkmgr.import_tf(self.gpu_config)
self.tf_sess = gpufmkmgr.get_tf_session()
self.keras = gpufmkmgr.import_keras()
self.keras_contrib = gpufmkmgr.import_keras_contrib()
self.onInitialize(**in_options)
if self.debug or self.batch_size == 0:
self.batch_size = 1
if self.is_training_mode:
if self.generator_list is None:
raise Exception('You didnt set_training_data_generators()')
else:
for i, generator in enumerate(self.generator_list):
if not isinstance(generator, SampleGeneratorBase):
raise Exception(
'training data generator is not subclass of SampleGeneratorBase'
)
if self.sample_for_preview is None:
self.sample_for_preview = self.generate_next_sample()
print("===== Model summary =====")
print("== Model name: " + self.get_model_name())
print("==")
print("== Current epoch: " + str(self.epoch))
print("==")
print("== Options:")
print("== |== batch_size : %s " % (self.batch_size))
print("== |== multi_gpu : %s " % (self.gpu_config.multi_gpu))
for key in self.options.keys():
print("== |== %s : %s" % (key, self.options[key]))
print("== Running on:")
if self.gpu_config.cpu_only:
print("== |== [CPU]")
else:
for idx in self.gpu_config.gpu_idxs:
print("== |== [%d : %s]" % (idx, gpufmkmgr.getDeviceName(idx)))
if not self.gpu_config.cpu_only and self.gpu_total_vram_gb == 2:
print("==")
print(
"== WARNING: You are using 2GB GPU. Result quality may be significantly decreased."
)
print(
"== If training does not start, close all programs and try again."
)
print(
"== Also you can disable Windows Aero Desktop to get extra free VRAM."
)
print("==")
print("=========================")