github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 04:52:13 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions
DeepFaceLab/models/ModelBase.py
Colombo 76ca79216e Upgraded to TF version 1.13.2 
Removed the wait at first launch for most graphics cards.

Increased speed of training by 10-20%, but you have to retrain all models from scratch.

SAEHD:

added option 'use float16'
	Experimental option. Reduces the model size by half.
	Increases the speed of training.
	Decreases the accuracy of the model.
	The model may collapse or not train.
	Model may not learn the mask in large resolutions.

true_face_training option is replaced by
"True face power". 0.0000 .. 1.0
Experimental option. Discriminates the result face to be more like the src face. Higher value - stronger discrimination.
Comparison - https://i.imgur.com/czScS9q.png
2020-01-25 21:58:19 +04:00

574 lines
24 KiB
Python
Raw Blame History

import colorsys
import inspect
import json
import operator
import os
import pickle
import shutil
import tempfile
import time
from pathlib import Path
import cv2
import numpy as np
from core import imagelib
from core.interact import interact as io
from core.leras import nn
from samplelib import SampleGeneratorBase
from core import pathex
from core.cv2ex import *
class ModelBase(object):
def __init__(self, is_training=False,
saved_models_path=None,
training_data_src_path=None,
training_data_dst_path=None,
pretraining_data_path=None,
pretrained_model_path=None,
no_preview=False,
force_model_name=None,
force_gpu_idxs=None,
cpu_only=False,
debug=False,
**kwargs):
self.is_training = is_training
self.saved_models_path = saved_models_path
self.training_data_src_path = training_data_src_path
self.training_data_dst_path = training_data_dst_path
self.pretraining_data_path = pretraining_data_path
self.pretrained_model_path = pretrained_model_path
self.no_preview = no_preview
self.debug = debug
self.model_class_name = model_class_name = Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
if force_model_name is not None:
self.model_name = force_model_name
else:
while True:
# gather all model dat files
saved_models_names = []
for filepath in pathex.get_file_paths(saved_models_path):
filepath_name = filepath.name
if filepath_name.endswith(f'{model_class_name}_data.dat'):
saved_models_names += [ (filepath_name.split('_')[0], os.path.getmtime(filepath)) ]
# sort by modified datetime
saved_models_names = sorted(saved_models_names, key=operator.itemgetter(1), reverse=True )
saved_models_names = [ x[0] for x in saved_models_names ]
if len(saved_models_names) != 0:
io.log_info ("Choose one of saved models, or enter a name to create a new model.")
io.log_info ("[r] : rename")
io.log_info ("[d] : delete")
io.log_info ("")
for i, model_name in enumerate(saved_models_names):
s = f"[{i}] : {model_name} "
if i == 0:
s += "- latest"
io.log_info (s)
inp = io.input_str(f"", "0", show_default_value=False )
model_idx = -1
try:
model_idx = np.clip ( int(inp), 0, len(saved_models_names)-1 )
except:
pass
if model_idx == -1:
if len(inp) == 1:
is_rename = inp[0] == 'r'
is_delete = inp[0] == 'd'
if is_rename or is_delete:
if len(saved_models_names) != 0:
if is_rename:
name = io.input_str(f"Enter the name of the model you want to rename")
elif is_delete:
name = io.input_str(f"Enter the name of the model you want to delete")
if name in saved_models_names:
if is_rename:
new_model_name = io.input_str(f"Enter new name of the model")
for filepath in pathex.get_file_paths(saved_models_path):
filepath_name = filepath.name
model_filename, remain_filename = filepath_name.split('_', 1)
if model_filename == name:
if is_rename:
new_filepath = filepath.parent / ( new_model_name + '_' + remain_filename )
filepath.rename (new_filepath)
elif is_delete:
filepath.unlink()
continue
self.model_name = inp
else:
self.model_name = saved_models_names[model_idx]
else:
self.model_name = io.input_str(f"No saved models found. Enter a name of a new model", "noname")
break
self.model_name = self.model_name + '_' + self.model_class_name
self.iter = 0
self.options = {}
self.loss_history = []
self.sample_for_preview = None
self.choosed_gpu_indexes = None
model_data = {}
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
if self.model_data_path.exists():
io.log_info (f"Loading {self.model_name} model...")
model_data = pickle.loads ( self.model_data_path.read_bytes() )
self.iter = model_data.get('iter',0)
if self.iter != 0:
self.options = model_data['options']
self.loss_history = model_data.get('loss_history', [])
self.sample_for_preview = model_data.get('sample_for_preview', None)
self.choosed_gpu_indexes = model_data.get('choosed_gpu_indexes', None)
if self.is_first_run():
io.log_info ("\nModel first run.")
self.device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_best_multi_gpu=True)) \
if not cpu_only else nn.DeviceConfig.CPU()
nn.initialize(self.device_config)
####
self.default_options_path = saved_models_path / f'{self.model_class_name}_default_options.dat'
self.default_options = {}
if self.default_options_path.exists():
try:
self.default_options = pickle.loads ( self.default_options_path.read_bytes() )
except:
pass
self.choose_preview_history = False
self.batch_size = self.load_or_def_option('batch_size', 1)
#####
io.input_skip_pending()
self.on_initialize_options()
if self.is_first_run():
# save as default options only for first run model initialize
self.default_options_path.write_bytes( pickle.dumps (self.options) )
self.autobackup = self.options.get('autobackup', False)
self.write_preview_history = self.options.get('write_preview_history', False)
self.target_iter = self.options.get('target_iter',0)
self.random_flip = self.options.get('random_flip',True)
self.on_initialize()
self.options['batch_size'] = self.batch_size
if self.is_training:
self.preview_history_path = self.saved_models_path / ( f'{self.get_model_name()}_history' )
self.autobackups_path = self.saved_models_path / ( f'{self.get_model_name()}_autobackups' )
if self.autobackup:
self.autobackup_current_hour = time.localtime().tm_hour
if not self.autobackups_path.exists():
self.autobackups_path.mkdir(exist_ok=True)
if self.write_preview_history or io.is_colab():
if not self.preview_history_path.exists():
self.preview_history_path.mkdir(exist_ok=True)
else:
if self.iter == 0:
for filename in pathex.get_image_paths(self.preview_history_path):
Path(filename).unlink()
if self.generator_list is None:
raise ValueError( 'You didnt set_training_data_generators()')
else:
for i, generator in enumerate(self.generator_list):
if not isinstance(generator, SampleGeneratorBase):
raise ValueError('training data generator is not subclass of SampleGeneratorBase')
if self.sample_for_preview is None or self.choose_preview_history:
if self.choose_preview_history and io.is_support_windows():
io.log_info ("Choose image for the preview history. [p] - next. [enter] - confirm.")
wnd_name = "[p] - next. [enter] - confirm."
io.named_window(wnd_name)
io.capture_keys(wnd_name)
choosed = False
while not choosed:
self.sample_for_preview = self.generate_next_samples()
preview = self.get_static_preview()
io.show_image( wnd_name, (preview*255).astype(np.uint8) )
while True:
key_events = io.get_key_events(wnd_name)
key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False)
if key == ord('\n') or key == ord('\r'):
choosed = True
break
elif key == ord('p'):
break
try:
io.process_messages(0.1)
except KeyboardInterrupt:
choosed = True
io.destroy_window(wnd_name)
else:
self.sample_for_preview = self.generate_next_samples()
try:
self.get_static_preview()
except:
self.sample_for_preview = self.generate_next_samples()
self.last_sample = self.sample_for_preview
io.log_info( self.get_summary_text() )
def load_or_def_option(self, name, def_value):
options_val = self.options.get(name, None)
if options_val is not None:
return options_val
def_opt_val = self.default_options.get(name, None)
if def_opt_val is not None:
return def_opt_val
return def_value
def ask_override(self):
return self.is_training and self.iter != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 5 if io.is_colab() else 2 )
def ask_enable_autobackup(self):
default_autobackup = self.options['autobackup'] = self.load_or_def_option('autobackup', False)
self.options['autobackup'] = io.input_bool(f"Enable autobackup", default_autobackup, help_message="Autobackup model files with preview every hour for last 15 hours. Latest backup located in model/<>_autobackups/01")
def ask_write_preview_history(self):
default_write_preview_history = self.load_or_def_option('write_preview_history', False)
self.options['write_preview_history'] = io.input_bool(f"Write preview history", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
if self.options['write_preview_history']:
if io.is_support_windows():
self.choose_preview_history = io.input_bool("Choose image for the preview history", False)
elif io.is_colab():
self.choose_preview_history = io.input_bool("Randomly choose new image for preview history", False, help_message="Preview image history will stay stuck with old faces if you reuse the same model on different celebs. Choose no unless you are changing src/dst to a new person")
def ask_target_iter(self):
default_target_iter = self.load_or_def_option('target_iter', 0)
self.options['target_iter'] = max(0, io.input_int("Target iteration", default_target_iter))
def ask_random_flip(self):
default_random_flip = self.load_or_def_option('random_flip', True)
self.options['random_flip'] = io.input_bool("Flip faces randomly", default_random_flip, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
def ask_batch_size(self, suggest_batch_size=None):
default_batch_size = self.load_or_def_option('batch_size', suggest_batch_size or self.batch_size)
self.batch_size = max(0, io.input_int("Batch_size", default_batch_size, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
#overridable
def on_initialize_options(self):
pass
#overridable
def on_initialize(self):
'''
initialize your models
store and retrieve your model options in self.options['']
check example
'''
pass
#overridable
def onSave(self):
#save your models here
pass
#overridable
def onTrainOneIter(self, sample, generator_list):
#train your models here
#return array of losses
return ( ('loss_src', 0), ('loss_dst', 0) )
#overridable
def onGetPreview(self, sample):
#you can return multiple previews
#return [ ('preview_name',preview_rgb), ... ]
return []
#overridable if you want model name differs from folder name
def get_model_name(self):
return self.model_name
#overridable , return [ [model, filename],... ] list
def get_model_filename_list(self):
return []
#overridable
def get_MergerConfig(self):
#return predictor_func, predictor_input_shape, MergerConfig() for the model
raise NotImplementedError
def get_pretraining_data_path(self):
return self.pretraining_data_path
def get_target_iter(self):
return self.target_iter
def is_reached_iter_goal(self):
return self.target_iter != 0 and self.iter >= self.target_iter
def get_previews(self):
return self.onGetPreview ( self.last_sample )
def get_static_preview(self):
return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr
def save(self):
summary_path = self.get_strpath_storage_for_file('summary.txt')
Path( summary_path ).write_text( self.get_summary_text() )
self.onSave()
model_data = {
'iter': self.iter,
'options': self.options,
'loss_history': self.loss_history,
'sample_for_preview' : self.sample_for_preview,
'choosed_gpu_indexes' : self.choosed_gpu_indexes,
}
pathex.write_bytes_safe (self.model_data_path, pickle.dumps(model_data) )
if self.autobackup:
bckp_filename_list = [ self.get_strpath_storage_for_file(filename) for _, filename in self.get_model_filename_list() ]
bckp_filename_list += [ str(summary_path), str(self.model_data_path) ]
current_hour = time.localtime().tm_hour
if self.autobackup_current_hour != current_hour:
self.autobackup_current_hour = current_hour
for i in range(15,0,-1):
idx_str = '%.2d' % i
next_idx_str = '%.2d' % (i+1)
idx_backup_path = self.autobackups_path / idx_str
next_idx_packup_path = self.autobackups_path / next_idx_str
if idx_backup_path.exists():
if i == 15:
pathex.delete_all_files(idx_backup_path)
else:
next_idx_packup_path.mkdir(exist_ok=True)
pathex.move_all_files (idx_backup_path, next_idx_packup_path)
if i == 1:
idx_backup_path.mkdir(exist_ok=True)
for filename in bckp_filename_list:
shutil.copy ( str(filename), str(idx_backup_path / Path(filename).name) )
previews = self.get_previews()
plist = []
for i in range(len(previews)):
name, bgr = previews[i]
plist += [ (bgr, idx_backup_path / ( ('preview_%s.jpg') % (name)) ) ]
for preview, filepath in plist:
preview_lh = ModelBase.get_loss_history_preview(self.loss_history, self.iter, preview.shape[1], preview.shape[2])
img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)
cv2_imwrite (filepath, img )
def debug_one_iter(self):
images = []
for generator in self.generator_list:
for i,batch in enumerate(next(generator)):
if len(batch.shape) == 4:
images.append( batch[0] )
return imagelib.equalize_and_stack_square (images)
def generate_next_samples(self):
self.last_sample = sample = [ generator.generate_next() for generator in self.generator_list]
return sample
def train_one_iter(self):
iter_time = time.time()
losses = self.onTrainOneIter()
iter_time = time.time() - iter_time
self.loss_history.append ( [float(loss[1]) for loss in losses] )
if self.iter % 10 == 0:
plist = []
if io.is_colab():
previews = self.get_previews()
for i in range(len(previews)):
name, bgr = previews[i]
plist += [ (bgr, self.get_strpath_storage_for_file('preview_%s.jpg' % (name) ) ) ]
if self.write_preview_history:
plist += [ (self.get_static_preview(), str (self.preview_history_path / ('%.6d.jpg' % (self.iter))) ) ]
for preview, filepath in plist:
preview_lh = ModelBase.get_loss_history_preview(self.loss_history, self.iter, preview.shape[1], preview.shape[2])
img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)
cv2_imwrite (filepath, img )
self.iter += 1
return self.iter, iter_time
def pass_one_iter(self):
self.generate_next_samples()
def finalize(self):
nn.tf_close_session()
def is_first_run(self):
return self.iter == 0
def is_debug(self):
return self.debug
def set_batch_size(self, batch_size):
self.batch_size = batch_size
def get_batch_size(self):
return self.batch_size
def get_iter(self):
return self.iter
def set_iter(self, iter):
self.iter = iter
self.loss_history = self.loss_history[:iter]
def get_loss_history(self):
return self.loss_history
def set_training_data_generators (self, generator_list):
self.generator_list = generator_list
def get_training_data_generators (self):
return self.generator_list
def get_model_root_path(self):
return self.saved_models_path
def get_strpath_storage_for_file(self, filename):
return str( self.saved_models_path / ( self.get_model_name() + '_' + filename) )
def get_summary_text(self):
###Generate text summary of model hyperparameters
#Find the longest key name and value string. Used as column widths.
width_name = max([len(k) for k in self.options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration"
width_value = max([len(str(x)) for x in self.options.values()] + [len(str(self.get_iter())), len(self.get_model_name())]) + 1 # Single space buffer to right edge
if len(self.device_config.devices) != 0: #Check length of GPU names
width_value = max([len(device.name)+1 for device in self.device_config.devices] + [width_value])
width_total = width_name + width_value + 2 #Plus 2 for ": "
summary_text = []
summary_text += [f'=={" Model Summary ":=^{width_total}}=='] # Model/status summary
summary_text += [f'=={" "*width_total}==']
summary_text += [f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='] # Name
summary_text += [f'=={" "*width_total}==']
summary_text += [f'=={"Current iteration": >{width_name}}: {str(self.get_iter()): <{width_value}}=='] # Iter
summary_text += [f'=={" "*width_total}==']
summary_text += [f'=={" Model Options ":-^{width_total}}=='] # Model options
summary_text += [f'=={" "*width_total}==']
for key in self.options.keys():
summary_text += [f'=={key: >{width_name}}: {str(self.options[key]): <{width_value}}=='] # self.options key/value pairs
summary_text += [f'=={" "*width_total}==']
summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info
summary_text += [f'=={" "*width_total}==']
if len(self.device_config.devices) == 0:
summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only
else:
for device in self.device_config.devices:
summary_text += [f'=={"Device index": >{width_name}}: {device.index: <{width_value}}=='] # GPU hardware device index
summary_text += [f'=={"Name": >{width_name}}: {device.name: <{width_value}}=='] # GPU name
vram_str = f'{device.total_mem_gb:.2f}GB' # GPU VRAM - Formated as #.## (or ##.##)
summary_text += [f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}==']
summary_text += [f'=={" "*width_total}==']
summary_text += [f'=={"="*width_total}==']
summary_text = "\n".join (summary_text)
return summary_text
@staticmethod
def get_loss_history_preview(loss_history, iter, w, c):
loss_history = np.array (loss_history.copy())
lh_height = 100
lh_img = np.ones ( (lh_height,w,c) ) * 0.1
if len(loss_history) != 0:
loss_count = len(loss_history[0])
lh_len = len(loss_history)
l_per_col = lh_len / w
plist_max = [ [ max (0.0, loss_history[int(col*l_per_col)][p],
*[ loss_history[i_ab][p]
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
]
)
for p in range(loss_count)
]
for col in range(w)
]
plist_min = [ [ min (plist_max[col][p], loss_history[int(col*l_per_col)][p],
*[ loss_history[i_ab][p]
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
]
)
for p in range(loss_count)
]
for col in range(w)
]
plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2
for col in range(0, w):
for p in range(0,loss_count):
point_color = [1.0]*c
point_color[0:3] = colorsys.hsv_to_rgb ( p * (1.0/loss_count), 1.0, 1.0 )
ph_max = int ( (plist_max[col][p] / plist_abs_max) * (lh_height-1) )
ph_max = np.clip( ph_max, 0, lh_height-1 )
ph_min = int ( (plist_min[col][p] / plist_abs_max) * (lh_height-1) )
ph_min = np.clip( ph_min, 0, lh_height-1 )
for ph in range(ph_min, ph_max+1):
lh_img[ (lh_height-ph-1), col ] = point_color
lh_lines = 5
lh_line_height = (lh_height-1)/lh_lines
for i in range(0,lh_lines+1):
lh_img[ int(i*lh_line_height), : ] = (0.8,)*c
last_line_t = int((lh_lines-1)*lh_line_height)
last_line_b = int(lh_lines*lh_line_height)
lh_text = 'Iter: %d' % (iter) if iter != 0 else ''
lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image ( (last_line_b-last_line_t,w,c), lh_text, color=[0.8]*c )
return lh_img
Reference in a new issue View git blame Copy permalink