Merge remote-tracking branch 'origin/config_files' into amp_test_config

This commit is contained in:
Jan 2021-12-05 17:57:19 +01:00
commit cbdf9f05c9
2 changed files with 15 additions and 8 deletions

View file

@ -13,6 +13,13 @@ from samplelib import *
from pathlib import Path from pathlib import Path
class QModel(ModelBase): class QModel(ModelBase):
#override
def on_initialize_options(self):
ask_override = False if self.read_from_conf else self.ask_override()
if self.is_first_run() or ask_override:
if (self.read_from_conf and not self.config_file_exists) or not self.read_from_conf:
self.ask_batch_size()
#override #override
def on_initialize(self): def on_initialize(self):
device_config = nn.getCurrentDeviceConfig() device_config = nn.getCurrentDeviceConfig()
@ -82,7 +89,7 @@ class QModel(ModelBase):
if self.is_training: if self.is_training:
# Adjust batch size for multiple GPU # Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) ) gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, 4 // gpu_count) bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu) self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU # Compute losses per GPU

View file

@ -20,7 +20,7 @@ class XSegModel(ModelBase):
#override #override
def on_initialize_options(self): def on_initialize_options(self):
ask_override = self.ask_override() ask_override = False if self.read_from_conf else self.ask_override()
if not self.is_first_run() and ask_override: if not self.is_first_run() and ask_override:
if io.input_bool(f"Restart training?", False, help_message="Reset model weights and start training from scratch."): if io.input_bool(f"Restart training?", False, help_message="Reset model weights and start training from scratch."):
@ -30,9 +30,11 @@ class XSegModel(ModelBase):
default_pretrain = self.options['pretrain'] = self.load_or_def_option('pretrain', False) default_pretrain = self.options['pretrain'] = self.load_or_def_option('pretrain', False)
if self.is_first_run(): if self.is_first_run():
if (self.read_from_conf and not self.config_file_exists) or not self.read_from_conf:
self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf','head'], help_message="Half / mid face / full face / whole face / head. Choose the same as your deepfake model.").lower() self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf','head'], help_message="Half / mid face / full face / whole face / head. Choose the same as your deepfake model.").lower()
if self.is_first_run() or ask_override: if self.is_first_run() or ask_override:
if (self.read_from_conf and not self.config_file_exists) or not self.read_from_conf:
self.ask_batch_size(4, range=[2,16]) self.ask_batch_size(4, range=[2,16])
self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain) self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain)
@ -53,14 +55,12 @@ class XSegModel(ModelBase):
self.resolution = resolution = 256 self.resolution = resolution = 256
self.face_type = {'h' : FaceType.HALF, self.face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL, 'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL, 'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE, 'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[ self.options['face_type'] ] 'head' : FaceType.HEAD}[ self.options['face_type'] ]
place_model_on_cpu = len(devices) == 0 place_model_on_cpu = len(devices) == 0
models_opt_device = '/CPU:0' if place_model_on_cpu else nn.tf_default_device_name models_opt_device = '/CPU:0' if place_model_on_cpu else nn.tf_default_device_name