added tf_cpu_mode option for funit models

This commit is contained in:
Colombo 2019-09-21 09:44:34 +04:00
parent d9d10f91c2
commit 2a3b3f0021
3 changed files with 20 additions and 5 deletions

View file

@ -33,6 +33,12 @@ class FUNITModel(ModelBase):
else: else:
self.options['face_type'] = self.options.get('face_type', default_face_type) self.options['face_type'] = self.options.get('face_type', default_face_type)
if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1)
self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.")
else:
self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1)
#override #override
def onInitialize(self, batch_size=-1, **in_options): def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals()) exec(nnlib.code_import_all, locals(), globals())
@ -58,7 +64,8 @@ class FUNITModel(ModelBase):
num_classes=person_id_max_count, num_classes=person_id_max_count,
subpixel_decoder=True, subpixel_decoder=True,
initialize_weights=self.is_first_run(), initialize_weights=self.is_first_run(),
is_training=self.is_training_mode is_training=self.is_training_mode,
tf_cpu_mode=self.options['optimizer_mode']-1
) )
if not self.is_first_run(): if not self.is_first_run():

View file

@ -32,6 +32,12 @@ class TrueFaceModel(ModelBase):
else: else:
self.options['face_type'] = self.options.get('face_type', default_face_type) self.options['face_type'] = self.options.get('face_type', default_face_type)
if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1)
self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.")
else:
self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1)
#override #override
def onInitialize(self, batch_size=-1, **in_options): def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals()) exec(nnlib.code_import_all, locals(), globals())
@ -55,7 +61,8 @@ class TrueFaceModel(ModelBase):
num_classes=2, num_classes=2,
subpixel_decoder=True, subpixel_decoder=True,
initialize_weights=self.is_first_run(), initialize_weights=self.is_first_run(),
is_training=self.is_training_mode is_training=self.is_training_mode,
tf_cpu_mode=self.options['optimizer_mode']-1
) )
if not self.is_first_run(): if not self.is_first_run():

View file

@ -30,7 +30,8 @@ class FUNIT(object):
load_weights_locally=False, load_weights_locally=False,
weights_file_root=None, weights_file_root=None,
is_training=True is_training=True,
tf_cpu_mode=0,
): ):
exec( nnlib.import_all(), locals(), globals() ) exec( nnlib.import_all(), locals(), globals() )
@ -47,8 +48,8 @@ class FUNIT(object):
self.dis = modelify ( FUNIT.DiscriminatorFlow(nf=dis_nf, n_res_blks=dis_res_blks, num_classes=num_classes) ) (Input(bgr_shape)) self.dis = modelify ( FUNIT.DiscriminatorFlow(nf=dis_nf, n_res_blks=dis_res_blks, num_classes=num_classes) ) (Input(bgr_shape))
self.G_opt = RMSprop(lr=0.0001, decay=0.0001, tf_cpu_mode=2 if 'tensorflow' in nnlib.active_DeviceConfig.backend else 0) self.G_opt = RMSprop(lr=0.0001, decay=0.0001, tf_cpu_mode=tf_cpu_mode)
self.D_opt = RMSprop(lr=0.0001, decay=0.0001, tf_cpu_mode=2 if 'tensorflow' in nnlib.active_DeviceConfig.backend else 0) self.D_opt = RMSprop(lr=0.0001, decay=0.0001, tf_cpu_mode=tf_cpu_mode)
xa = Input(bgr_shape, name="xa") xa = Input(bgr_shape, name="xa")
la = Input(label_shape, dtype="int32", name="la") la = Input(label_shape, dtype="int32", name="la")