mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 13:02:15 -07:00
now you can train models on multiple GPU's on same workspace without cloning any folders.
Model files names will be prefixed with GPU index if GPU choosed explicitly on train/convert start. if you leave GPU idx choice default, then best GPU idx will be choosed and model file names will not contain index prefix. It gives you possibility to train same fake with various models or options on multiple GPUs. H64 and H128: now you can choose 'Lighter autoencoder'. It is same as vram gb <= 4 before this update. added archived_models.zip contains old experiments RecycleGAN: archived devicelib: if your system has no NVML installed (some old cards), then it will work with gpu_idx=0 as 'Generic GeForce GPU' with 2GB vram. refactorings
This commit is contained in:
iperov
parent
e2f4677987
commit
1f2b1481ef
9 changed files with 180 additions and 479 deletions
22
main.py
22
main.py
|
|
@ -65,16 +65,6 @@ if __name__ == "__main__":
|
|||
sort_parser.set_defaults (func=process_sort)
|
||||
|
||||
def process_train(arguments):
|
||||
|
||||
if 'DFL_TARGET_EPOCH' in os.environ.keys():
|
||||
arguments.session_target_epoch = int ( os.environ['DFL_TARGET_EPOCH'] )
|
||||
|
||||
if 'DFL_BATCH_SIZE' in os.environ.keys():
|
||||
arguments.batch_size = int ( os.environ['DFL_BATCH_SIZE'] )
|
||||
|
||||
if 'DFL_WORST_GPU' in os.environ.keys():
|
||||
arguments.choose_worst_gpu = True
|
||||
|
||||
from mainscripts import Trainer
|
||||
Trainer.main (
|
||||
training_data_src_dir=arguments.training_data_src_dir,
|
||||
|
|
@ -83,9 +73,7 @@ if __name__ == "__main__":
|
|||
model_name=arguments.model_name,
|
||||
debug = arguments.debug,
|
||||
#**options
|
||||
choose_worst_gpu = arguments.choose_worst_gpu,
|
||||
force_best_gpu_idx = arguments.force_best_gpu_idx,
|
||||
force_gpu_idxs = arguments.force_gpu_idxs,
|
||||
force_gpu_idx = arguments.force_gpu_idx,
|
||||
cpu_only = arguments.cpu_only
|
||||
)
|
||||
|
||||
|
|
@ -96,9 +84,7 @@ if __name__ == "__main__":
|
|||
train_parser.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model")
|
||||
train_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.")
|
||||
train_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
|
||||
train_parser.add_argument('--force-gpu-idxs', type=str, dest="force_gpu_idxs", default=None, help="Override final GPU idxs. Example: 0,1,2.")
|
||||
train_parser.add_argument('--choose-worst-gpu', action="store_true", dest="choose_worst_gpu", default=False, help="Choose worst GPU instead of best. Environment variable to force True: DFL_WORST_GPU")
|
||||
train_parser.add_argument('--force-best-gpu-idx', type=int, dest="force_best_gpu_idx", default=-1, help="Force to choose this GPU idx as best(worst).")
|
||||
train_parser.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.")
|
||||
|
||||
train_parser.set_defaults (func=process_train)
|
||||
|
||||
|
|
@ -111,7 +97,7 @@ if __name__ == "__main__":
|
|||
model_dir=arguments.model_dir,
|
||||
model_name=arguments.model_name,
|
||||
debug = arguments.debug,
|
||||
force_best_gpu_idx = arguments.force_best_gpu_idx,
|
||||
force_gpu_idx = arguments.force_gpu_idx,
|
||||
cpu_only = arguments.cpu_only
|
||||
)
|
||||
|
||||
|
|
@ -122,7 +108,7 @@ if __name__ == "__main__":
|
|||
convert_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
|
||||
convert_parser.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model")
|
||||
convert_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.")
|
||||
convert_parser.add_argument('--force-best-gpu-idx', type=int, dest="force_best_gpu_idx", default=-1, help="Force to choose this GPU idx as best.")
|
||||
convert_parser.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.")
|
||||
convert_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.")
|
||||
|
||||
convert_parser.set_defaults(func=process_convert)
|
||||
|
|
|
|||
|
|
@ -68,7 +68,7 @@ class ExtractSubprocessor(SubprocessorBase):
|
|||
if not multi_gpu or len(devices) == 0:
|
||||
devices = [nnlib.device.getBestDeviceIdx()]
|
||||
|
||||
if len(devices) == 0 or devices[0] == -1:
|
||||
if len(devices) == 0:
|
||||
devices = [0]
|
||||
|
||||
devices = [ (idx, nnlib.device.getDeviceName(idx), nnlib.device.getDeviceVRAMTotalGb(idx) ) for idx in devices]
|
||||
|
|
@ -263,7 +263,7 @@ class ExtractSubprocessor(SubprocessorBase):
|
|||
|
||||
self.e = None
|
||||
|
||||
device_config = nnlib.DeviceConfig ( cpu_only=self.cpu_only, force_best_gpu_idx=self.device_idx, allow_growth=True)
|
||||
device_config = nnlib.DeviceConfig ( cpu_only=self.cpu_only, force_gpu_idx=self.device_idx, allow_growth=True)
|
||||
if self.type == 'rects':
|
||||
if self.detector is not None:
|
||||
if self.detector == 'mt':
|
||||
|
|
|
|||
|
|
@ -18,7 +18,18 @@ You can implement your own model. Check examples.
|
|||
class ModelBase(object):
|
||||
|
||||
#DONT OVERRIDE
|
||||
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, debug = False, force_best_gpu_idx=-1, **in_options):
|
||||
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, debug = False, force_gpu_idx=-1, **in_options):
|
||||
|
||||
if force_gpu_idx == -1:
|
||||
idxs_names_list = nnlib.device.getAllDevicesIdxsWithNamesList()
|
||||
if len(idxs_names_list) > 1:
|
||||
print ("You have multi GPUs in a system: ")
|
||||
for idx, name in idxs_names_list:
|
||||
print ("[%d] : %s" % (idx, name) )
|
||||
|
||||
force_gpu_idx = input_int("Which GPU idx to choose? ( skip: best GPU ) : ", -1, [ x[0] for x in idxs_names_list] )
|
||||
self.force_gpu_idx = force_gpu_idx
|
||||
|
||||
print ("Loading model...")
|
||||
self.model_path = model_path
|
||||
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
|
||||
|
|
@ -35,7 +46,7 @@ class ModelBase(object):
|
|||
self.debug = debug
|
||||
self.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None)
|
||||
|
||||
self.supress_std_once = ('TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1')
|
||||
self.supress_std_once = os.environ.get('TF_SUPPRESS_STD', '0') == '1'
|
||||
|
||||
self.epoch = 0
|
||||
self.options = {}
|
||||
|
|
@ -54,15 +65,6 @@ class ModelBase(object):
|
|||
if self.epoch == 0:
|
||||
print ("\nModel first run. Enter model options as default for each run.")
|
||||
|
||||
if (self.epoch == 0 or ask_override) and (force_best_gpu_idx == -1):
|
||||
idxs_names_list = nnlib.device.getAllDevicesIdxsWithNamesList()
|
||||
if len(idxs_names_list) > 1:
|
||||
print ("You have multi GPUs in a system: ")
|
||||
for idx, name in idxs_names_list:
|
||||
print ("[%d] : %s" % (idx, name) )
|
||||
|
||||
force_best_gpu_idx = input_int("Which GPU idx to choose? ( skip: system choice ) : ", -1)
|
||||
|
||||
if self.epoch == 0 or ask_override:
|
||||
default_write_preview_history = False if self.epoch == 0 else self.options.get('write_preview_history',False)
|
||||
self.options['write_preview_history'] = input_bool("Write preview history? (y/n ?:help skip:n/default) : ", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
|
||||
|
|
@ -119,14 +121,9 @@ class ModelBase(object):
|
|||
|
||||
self.onInitializeOptions(self.epoch == 0, ask_override)
|
||||
|
||||
nnlib.import_all ( nnlib.DeviceConfig(allow_growth=False, force_best_gpu_idx=force_best_gpu_idx, **in_options) )
|
||||
nnlib.import_all ( nnlib.DeviceConfig(allow_growth=False, force_gpu_idx=self.force_gpu_idx, **in_options) )
|
||||
self.device_config = nnlib.active_DeviceConfig
|
||||
|
||||
if self.epoch == 0:
|
||||
self.created_vram_gb = self.options['created_vram_gb'] = self.device_config.gpu_total_vram_gb
|
||||
else:
|
||||
self.created_vram_gb = self.options['created_vram_gb'] = self.options.get('created_vram_gb',self.device_config.gpu_total_vram_gb)
|
||||
|
||||
self.onInitialize(**in_options)
|
||||
|
||||
self.options['batch_size'] = self.batch_size
|
||||
|
|
@ -136,7 +133,10 @@ class ModelBase(object):
|
|||
|
||||
if self.is_training_mode:
|
||||
if self.write_preview_history:
|
||||
if self.force_gpu_idx == -1:
|
||||
self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) )
|
||||
else:
|
||||
self.preview_history_path = self.model_path / ( '%d_%s_history' % (self.force_gpu_idx, self.get_model_name()) )
|
||||
|
||||
if not self.preview_history_path.exists():
|
||||
self.preview_history_path.mkdir(exist_ok=True)
|
||||
|
|
@ -174,7 +174,7 @@ class ModelBase(object):
|
|||
for idx in self.device_config.gpu_idxs:
|
||||
print ("== |== [%d : %s]" % (idx, nnlib.device.getDeviceName(idx)) )
|
||||
|
||||
if not self.device_config.cpu_only and self.device_config.gpu_total_vram_gb == 2:
|
||||
if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[0] == 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.")
|
||||
|
|
@ -367,7 +367,10 @@ class ModelBase(object):
|
|||
return self.generator_list
|
||||
|
||||
def get_strpath_storage_for_file(self, filename):
|
||||
return str( self.model_path / (self.get_model_name() + '_' + filename) )
|
||||
if self.force_gpu_idx == -1:
|
||||
return str( self.model_path / ( self.get_model_name() + '_' + filename) )
|
||||
else:
|
||||
return str( self.model_path / ( str(self.force_gpu_idx) + '_' + self.get_model_name() + '_' + filename) )
|
||||
|
||||
def set_vram_batch_requirements (self, d):
|
||||
#example d = {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48}
|
||||
|
|
@ -379,7 +382,7 @@ class ModelBase(object):
|
|||
else:
|
||||
if self.batch_size == 0:
|
||||
for x in keys:
|
||||
if self.device_config.gpu_total_vram_gb <= x:
|
||||
if self.device_config.gpu_vram_gb[0] <= x:
|
||||
self.batch_size = d[x]
|
||||
break
|
||||
|
||||
|
|
|
|||
|
|
@ -11,12 +11,22 @@ class Model(ModelBase):
|
|||
decoder_srcH5 = 'decoder_src.h5'
|
||||
decoder_dstH5 = 'decoder_dst.h5'
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run:
|
||||
self.options['lighter_ae'] = input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
|
||||
else:
|
||||
default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future
|
||||
if 'created_vram_gb' in self.options.keys():
|
||||
self.options.pop ('created_vram_gb')
|
||||
self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {2.5:2,3:2,4:2,4:4,5:8,6:12,7:16,8:16,9:24,10:24,11:32,12:32,13:48} )
|
||||
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.created_vram_gb)
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build( self.options['lighter_ae'] )
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoder_src.load_weights (self.get_strpath_storage_for_file(self.decoder_srcH5))
|
||||
|
|
@ -120,7 +130,7 @@ class Model(ModelBase):
|
|||
base_blur_mask_modifier=100,
|
||||
**in_options)
|
||||
|
||||
def Build(self, created_vram_gb):
|
||||
def Build(self, lighter_ae):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
bgr_shape = (128, 128, 3)
|
||||
|
|
@ -139,7 +149,7 @@ class Model(ModelBase):
|
|||
def Encoder(input_shape):
|
||||
input_layer = Input(input_shape)
|
||||
x = input_layer
|
||||
if created_vram_gb >= 5:
|
||||
if not lighter_ae:
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
|
|
@ -161,7 +171,7 @@ class Model(ModelBase):
|
|||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
if created_vram_gb >= 5:
|
||||
if not lighter_ae:
|
||||
input_ = Input(shape=(16, 16, 512))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
|
|
|
|||
|
|
@ -12,12 +12,24 @@ class Model(ModelBase):
|
|||
decoder_srcH5 = 'decoder_src.h5'
|
||||
decoder_dstH5 = 'decoder_dst.h5'
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run:
|
||||
self.options['lighter_ae'] = input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
|
||||
else:
|
||||
default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future
|
||||
if 'created_vram_gb' in self.options.keys():
|
||||
self.options.pop ('created_vram_gb')
|
||||
self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {1.5:2,2:2,3:8,4:16,5:24,6:32,7:40,8:48} )
|
||||
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.created_vram_gb)
|
||||
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.options['lighter_ae'])
|
||||
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoder_src.load_weights (self.get_strpath_storage_for_file(self.decoder_srcH5))
|
||||
|
|
@ -122,7 +134,7 @@ class Model(ModelBase):
|
|||
base_blur_mask_modifier=100,
|
||||
**in_options)
|
||||
|
||||
def Build(self, created_vram_gb):
|
||||
def Build(self, lighter_ae):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
bgr_shape = (64, 64, 3)
|
||||
|
|
@ -141,7 +153,7 @@ class Model(ModelBase):
|
|||
def Encoder(input_shape):
|
||||
input_layer = Input(input_shape)
|
||||
x = input_layer
|
||||
if created_vram_gb >= 4:
|
||||
if not lighter_ae:
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
|
|
@ -162,7 +174,7 @@ class Model(ModelBase):
|
|||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
if created_vram_gb >= 4:
|
||||
if not lighter_ae:
|
||||
input_ = Input(shape=(8, 8, 512))
|
||||
x = input_
|
||||
|
||||
|
|
|
|||
|
|
@ -1,250 +0,0 @@
|
|||
from models import ModelBase
|
||||
import numpy as np
|
||||
import cv2
|
||||
from mathlib import get_power_of_two
|
||||
from nnlib import nnlib
|
||||
|
||||
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
GAH5 = 'GA.h5'
|
||||
PAH5 = 'PA.h5'
|
||||
DAH5 = 'DA.h5'
|
||||
GBH5 = 'GB.h5'
|
||||
DBH5 = 'DB.h5'
|
||||
PBH5 = 'PB.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, batch_size=-1, **in_options):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
created_batch_size = self.get_batch_size()
|
||||
if self.epoch == 0:
|
||||
#first run
|
||||
|
||||
|
||||
|
||||
try:
|
||||
created_resolution = int ( input ("Resolution (default:64, valid: 64,128,256) : ") )
|
||||
except:
|
||||
created_resolution = 64
|
||||
|
||||
if created_resolution not in [64,128,256]:
|
||||
created_resolution = 64
|
||||
|
||||
try:
|
||||
created_batch_size = int ( input ("Batch_size (minimum/default - 10) : ") )
|
||||
except:
|
||||
created_batch_size = 10
|
||||
created_batch_size = max(created_batch_size,1)
|
||||
|
||||
print ("Done. If training won't start, decrease resolution")
|
||||
|
||||
self.options['created_resolution'] = created_resolution
|
||||
self.options['created_batch_size'] = created_batch_size
|
||||
self.created_vram_gb = self.device_config.gpu_total_vram_gb
|
||||
else:
|
||||
#not first run
|
||||
if 'created_batch_size' in self.options.keys():
|
||||
created_batch_size = self.options['created_batch_size']
|
||||
else:
|
||||
raise Exception("Continue training, but created_batch_size not found.")
|
||||
|
||||
if 'created_resolution' in self.options.keys():
|
||||
created_resolution = self.options['created_resolution']
|
||||
else:
|
||||
raise Exception("Continue training, but created_resolution not found.")
|
||||
|
||||
resolution = created_resolution
|
||||
bgr_shape = (resolution, resolution, 3)
|
||||
ngf = 64
|
||||
npf = 64
|
||||
ndf = 64
|
||||
lambda_A = 10
|
||||
lambda_B = 10
|
||||
|
||||
self.set_batch_size(created_batch_size)
|
||||
|
||||
use_batch_norm = False #created_batch_size > 1
|
||||
self.GA = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
self.GB = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
#self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
#self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
|
||||
self.PA = modelify(UNetTemporalPredictor(bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=npf, use_dropout=True))([Input(bgr_shape), Input(bgr_shape)])
|
||||
self.PB = modelify(UNetTemporalPredictor(bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=npf, use_dropout=True))([Input(bgr_shape), Input(bgr_shape)])
|
||||
|
||||
self.DA = modelify(NLayerDiscriminator(use_batch_norm, ndf=ndf, n_layers=3) ) (Input(bgr_shape))
|
||||
self.DB = modelify(NLayerDiscriminator(use_batch_norm, ndf=ndf, n_layers=3) ) (Input(bgr_shape))
|
||||
|
||||
if not self.is_first_run():
|
||||
self.GA.load_weights (self.get_strpath_storage_for_file(self.GAH5))
|
||||
self.DA.load_weights (self.get_strpath_storage_for_file(self.DAH5))
|
||||
self.PA.load_weights (self.get_strpath_storage_for_file(self.PAH5))
|
||||
self.GB.load_weights (self.get_strpath_storage_for_file(self.GBH5))
|
||||
self.DB.load_weights (self.get_strpath_storage_for_file(self.DBH5))
|
||||
self.PB.load_weights (self.get_strpath_storage_for_file(self.PBH5))
|
||||
|
||||
real_A0 = Input(bgr_shape, name="real_A0")
|
||||
real_A1 = Input(bgr_shape, name="real_A1")
|
||||
real_A2 = Input(bgr_shape, name="real_A2")
|
||||
|
||||
real_B0 = Input(bgr_shape, name="real_B0")
|
||||
real_B1 = Input(bgr_shape, name="real_B1")
|
||||
real_B2 = Input(bgr_shape, name="real_B2")
|
||||
|
||||
DA_ones = K.ones ( K.int_shape(self.DA.outputs[0])[1:] )
|
||||
DA_zeros = K.zeros ( K.int_shape(self.DA.outputs[0])[1:] )
|
||||
DB_ones = K.ones ( K.int_shape(self.DB.outputs[0])[1:] )
|
||||
DB_zeros = K.zeros ( K.int_shape(self.DB.outputs[0])[1:] )
|
||||
|
||||
def CycleLoss (t1,t2):
|
||||
return K.mean(K.square(t1 - t2))
|
||||
|
||||
def RecurrentLOSS(t1,t2):
|
||||
return K.mean(K.square(t1 - t2))
|
||||
|
||||
def RecycleLOSS(t1,t2):
|
||||
return K.mean(K.square(t1 - t2))
|
||||
|
||||
fake_B0 = self.GA(real_A0)
|
||||
fake_B1 = self.GA(real_A1)
|
||||
|
||||
fake_A0 = self.GB(real_B0)
|
||||
fake_A1 = self.GB(real_B1)
|
||||
|
||||
#rec_FB0 = self.GA(fake_A0)
|
||||
#rec_FB1 = self.GA(fake_A1)
|
||||
|
||||
#rec_FA0 = self.GB(fake_B0)
|
||||
#rec_FA1 = self.GB(fake_B1)
|
||||
|
||||
pred_A2 = self.PA ( [real_A0, real_A1])
|
||||
pred_B2 = self.PB ( [real_B0, real_B1])
|
||||
rec_A2 = self.GB ( self.PB ( [fake_B0, fake_B1]) )
|
||||
rec_B2 = self.GA ( self.PA ( [fake_A0, fake_A1]))
|
||||
|
||||
loss_G = K.mean(K.square(self.DB(fake_B0) - DB_ones)) + \
|
||||
K.mean(K.square(self.DB(fake_B1) - DB_ones)) + \
|
||||
K.mean(K.square(self.DA(fake_A0) - DA_ones)) + \
|
||||
K.mean(K.square(self.DA(fake_A1) - DA_ones)) + \
|
||||
lambda_A * ( #CycleLoss(rec_FA0, real_A0) + \
|
||||
#CycleLoss(rec_FA1, real_A1) + \
|
||||
RecurrentLOSS(pred_A2, real_A2) + \
|
||||
RecycleLOSS(rec_A2, real_A2) ) + \
|
||||
lambda_B * ( #CycleLoss(rec_FB0, real_B0) + \
|
||||
#CycleLoss(rec_FB1, real_B1) + \
|
||||
RecurrentLOSS(pred_B2, real_B2) + \
|
||||
RecycleLOSS(rec_B2, real_B2) )
|
||||
|
||||
weights_G = self.GA.trainable_weights + self.GB.trainable_weights + self.PA.trainable_weights + self.PB.trainable_weights
|
||||
|
||||
self.G_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_G],
|
||||
Adam(lr=2e-4, beta_1=0.5, beta_2=0.999).get_updates(loss_G, weights_G) )
|
||||
|
||||
###########
|
||||
|
||||
loss_D_A0 = ( K.mean(K.square( self.DA(real_A0) - DA_ones)) + \
|
||||
K.mean(K.square( self.DA(fake_A0) - DA_zeros)) ) * 0.5
|
||||
|
||||
loss_D_A1 = ( K.mean(K.square( self.DA(real_A1) - DA_ones)) + \
|
||||
K.mean(K.square( self.DA(fake_A1) - DA_zeros)) ) * 0.5
|
||||
|
||||
loss_D_A = loss_D_A0 + loss_D_A1
|
||||
|
||||
self.DA_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_A],
|
||||
Adam(lr=2e-4, beta_1=0.5, beta_2=0.999).get_updates(loss_D_A, self.DA.trainable_weights) )
|
||||
|
||||
############
|
||||
|
||||
loss_D_B0 = ( K.mean(K.square( self.DB(real_B0) - DB_ones)) + \
|
||||
K.mean(K.square( self.DB(fake_B0) - DB_zeros)) ) * 0.5
|
||||
|
||||
loss_D_B1 = ( K.mean(K.square( self.DB(real_B1) - DB_ones)) + \
|
||||
K.mean(K.square( self.DB(fake_B1) - DB_zeros)) ) * 0.5
|
||||
|
||||
loss_D_B = loss_D_B0 + loss_D_B1
|
||||
|
||||
self.DB_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_B],
|
||||
Adam(lr=2e-4, beta_1=0.5, beta_2=0.999).get_updates(loss_D_B, self.DB.trainable_weights) )
|
||||
|
||||
############
|
||||
|
||||
|
||||
self.G_view = K.function([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 ])
|
||||
self.G_convert = K.function([real_B0],[fake_A0])
|
||||
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorImageTemporal(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
temporal_image_count=3,
|
||||
sample_process_options=SampleProcessor.Options(random_flip = False, normalize_tanh = True),
|
||||
output_sample_types=[ [f.SOURCE | f.MODE_BGR, resolution] ] ),
|
||||
|
||||
SampleGeneratorImageTemporal(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
temporal_image_count=3,
|
||||
sample_process_options=SampleProcessor.Options(random_flip = False, normalize_tanh = True),
|
||||
output_sample_types=[ [f.SOURCE | f.MODE_BGR, resolution] ] ),
|
||||
])
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( [[self.GA, self.get_strpath_storage_for_file(self.GAH5)],
|
||||
[self.GB, self.get_strpath_storage_for_file(self.GBH5)],
|
||||
[self.DA, self.get_strpath_storage_for_file(self.DAH5)],
|
||||
[self.DB, self.get_strpath_storage_for_file(self.DBH5)],
|
||||
[self.PA, self.get_strpath_storage_for_file(self.PAH5)],
|
||||
[self.PB, self.get_strpath_storage_for_file(self.PBH5)] ])
|
||||
|
||||
#override
|
||||
def onTrainOneEpoch(self, sample):
|
||||
source_src_0, source_src_1, source_src_2, = sample[0]
|
||||
source_dst_0, source_dst_1, source_dst_2, = sample[1]
|
||||
|
||||
feed = [source_src_0, source_src_1, source_src_2, source_dst_0, source_dst_1, source_dst_2]
|
||||
|
||||
loss_G, = self.G_train ( feed )
|
||||
loss_DA, = self.DA_train( feed )
|
||||
loss_DB, = self.DB_train( feed )
|
||||
#return ( ('G', loss_G), )
|
||||
return ( ('G', loss_G), ('DA', loss_DA), ('DB', loss_DB) )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A0 = sample[0][0]
|
||||
test_A1 = sample[0][1]
|
||||
test_A2 = sample[0][2]
|
||||
|
||||
test_B0 = sample[1][0]
|
||||
test_B1 = sample[1][1]
|
||||
test_B2 = sample[1][2]
|
||||
|
||||
G_view_result = self.G_view([test_A0, test_A1, test_A2, test_B0, test_B1, test_B2])
|
||||
|
||||
fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 = [ x[0] / 2 + 0.5 for x in G_view_result]
|
||||
test_A0, test_A1, test_A2, test_B0, test_B1, test_B2 = [ x[0] / 2 + 0.5 for x in [test_A0, test_A1, test_A2, test_B0, test_B1, test_B2] ]
|
||||
|
||||
|
||||
r = np.concatenate ((np.concatenate ( (test_A0, test_A1, test_A2, pred_A2, fake_B0, fake_B1, rec_A2), axis=1),
|
||||
np.concatenate ( (test_B0, test_B1, test_B2, pred_B2, fake_A0, fake_A1, rec_B2), axis=1)
|
||||
), axis=0)
|
||||
|
||||
return [ ('RecycleGAN, A0-A1-A2-PA2-FB0-FB1-RA2, B0-B1-B2-PB2-FA0-FA1-RB2, ', r ) ]
|
||||
|
||||
def predictor_func (self, face):
|
||||
x = self.G_convert ( [ np.expand_dims(face *2 - 1,0)] )[0]
|
||||
return x[0] / 2 + 0.5
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterImage
|
||||
return ConverterImage(self.predictor_func,
|
||||
predictor_input_size=self.options['created_resolution'],
|
||||
output_size=self.options['created_resolution'],
|
||||
**in_options)
|
||||
|
||||
|
|
@ -1 +0,0 @@
|
|||
from .Model import Model
|
||||
BIN
models/archived_models.zip
Normal file
BIN
models/archived_models.zip
Normal file
Binary file not shown.
|
|
@ -1,20 +1,26 @@
|
|||
from .pynvml import *
|
||||
|
||||
try:
|
||||
nvmlInit()
|
||||
hasNVML = True
|
||||
except:
|
||||
hasNVML = False
|
||||
|
||||
class devicelib:
|
||||
class Config():
|
||||
force_best_gpu_idx = -1
|
||||
force_gpu_idx = -1
|
||||
multi_gpu = False
|
||||
force_gpu_idxs = None
|
||||
choose_worst_gpu = False
|
||||
gpu_idxs = []
|
||||
gpu_names = []
|
||||
gpu_total_vram_gb = 0
|
||||
gpu_compute_caps = []
|
||||
gpu_vram_gb = []
|
||||
allow_growth = True
|
||||
use_fp16 = False
|
||||
cpu_only = False
|
||||
|
||||
def __init__ (self, force_best_gpu_idx = -1,
|
||||
def __init__ (self, force_gpu_idx = -1,
|
||||
multi_gpu = False,
|
||||
force_gpu_idxs = None,
|
||||
choose_worst_gpu = False,
|
||||
|
|
@ -27,7 +33,7 @@ class devicelib:
|
|||
if cpu_only:
|
||||
self.cpu_only = True
|
||||
else:
|
||||
self.force_best_gpu_idx = force_best_gpu_idx
|
||||
self.force_gpu_idx = force_gpu_idx
|
||||
self.multi_gpu = multi_gpu
|
||||
self.force_gpu_idxs = force_gpu_idxs
|
||||
self.choose_worst_gpu = choose_worst_gpu
|
||||
|
|
@ -35,19 +41,13 @@ class devicelib:
|
|||
|
||||
self.gpu_idxs = []
|
||||
|
||||
if not devicelib.hasNVML():
|
||||
self.gpu_idxs = [0]
|
||||
self.gpu_total_vram_gb = 2
|
||||
self.gpu_names += ['Generic GeForce GPU']
|
||||
self.gpu_compute_caps += [ 50 ]
|
||||
else:
|
||||
if force_gpu_idxs is not None:
|
||||
for idx in force_gpu_idxs.split(','):
|
||||
idx = int(idx)
|
||||
if devicelib.isValidDeviceIdx(idx):
|
||||
self.gpu_idxs.append(idx)
|
||||
else:
|
||||
gpu_idx = force_best_gpu_idx if (force_best_gpu_idx >= 0 and devicelib.isValidDeviceIdx(force_best_gpu_idx)) else devicelib.getBestDeviceIdx() if not choose_worst_gpu else devicelib.getWorstDeviceIdx()
|
||||
gpu_idx = force_gpu_idx if (force_gpu_idx >= 0 and devicelib.isValidDeviceIdx(force_gpu_idx)) else devicelib.getBestDeviceIdx() if not choose_worst_gpu else devicelib.getWorstDeviceIdx()
|
||||
if gpu_idx != -1:
|
||||
if self.multi_gpu:
|
||||
self.gpu_idxs = devicelib.getDeviceIdxsEqualModel( gpu_idx )
|
||||
|
|
@ -59,187 +59,128 @@ class devicelib:
|
|||
self.cpu_only = (len(self.gpu_idxs) == 0)
|
||||
|
||||
if not self.cpu_only:
|
||||
self.gpu_total_vram_gb = devicelib.getDeviceVRAMTotalGb ( self.gpu_idxs[0] )
|
||||
self.gpu_names = []
|
||||
self.gpu_compute_caps = []
|
||||
for gpu_idx in self.gpu_idxs:
|
||||
self.gpu_names += [devicelib.getDeviceName(gpu_idx)]
|
||||
self.gpu_compute_caps += [ devicelib.getDeviceComputeCapability ( gpu_idx ) ]
|
||||
|
||||
@staticmethod
|
||||
def hasNVML():
|
||||
try:
|
||||
nvmlInit()
|
||||
nvmlShutdown()
|
||||
except:
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def getDevicesWithAtLeastFreeMemory(freememsize):
|
||||
result = []
|
||||
try:
|
||||
nvmlInit()
|
||||
for i in range(0, nvmlDeviceGetCount() ):
|
||||
handle = nvmlDeviceGetHandleByIndex(i)
|
||||
memInfo = nvmlDeviceGetMemoryInfo( handle )
|
||||
if (memInfo.total - memInfo.used) >= freememsize:
|
||||
result.append (i)
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
self.gpu_vram_gb += [ devicelib.getDeviceVRAMTotalGb ( gpu_idx ) ]
|
||||
|
||||
@staticmethod
|
||||
def getDevicesWithAtLeastTotalMemoryGB(totalmemsize_gb):
|
||||
if not hasNVML and totalmemsize_gb <= 2:
|
||||
return [0]
|
||||
|
||||
result = []
|
||||
try:
|
||||
nvmlInit()
|
||||
for i in range(0, nvmlDeviceGetCount() ):
|
||||
for i in range(nvmlDeviceGetCount()):
|
||||
handle = nvmlDeviceGetHandleByIndex(i)
|
||||
memInfo = nvmlDeviceGetMemoryInfo( handle )
|
||||
if (memInfo.total) >= totalmemsize_gb*1024*1024*1024:
|
||||
result.append (i)
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def getAllDevicesIdxsList ():
|
||||
result = []
|
||||
try:
|
||||
nvmlInit()
|
||||
result = [ i for i in range(0, nvmlDeviceGetCount() ) ]
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
def getAllDevicesIdxsList():
|
||||
if not hasNVML:
|
||||
return [0]
|
||||
|
||||
return [ i for i in range(0, nvmlDeviceGetCount() ) ]
|
||||
|
||||
@staticmethod
|
||||
def getAllDevicesIdxsWithNamesList ():
|
||||
result = []
|
||||
try:
|
||||
nvmlInit()
|
||||
result = [ (i, nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)).decode() ) for i in range(0, nvmlDeviceGetCount() ) ]
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
def getAllDevicesIdxsWithNamesList():
|
||||
if not hasNVML:
|
||||
return [ (0, devicelib.getDeviceName(0) ) ]
|
||||
|
||||
return [ (i, nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)).decode() ) for i in range(nvmlDeviceGetCount() ) ]
|
||||
|
||||
@staticmethod
|
||||
def getDeviceVRAMFree (idx):
|
||||
result = 0
|
||||
try:
|
||||
nvmlInit()
|
||||
if not hasNVML:
|
||||
return 2
|
||||
|
||||
if idx < nvmlDeviceGetCount():
|
||||
handle = nvmlDeviceGetHandleByIndex(idx)
|
||||
memInfo = nvmlDeviceGetMemoryInfo( handle )
|
||||
result = (memInfo.total - memInfo.used)
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
memInfo = nvmlDeviceGetMemoryInfo( nvmlDeviceGetHandleByIndex(idx) )
|
||||
return memInfo.total - memInfo.used
|
||||
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def getDeviceVRAMTotalGb (idx):
|
||||
result = 2
|
||||
try:
|
||||
nvmlInit()
|
||||
if not hasNVML:
|
||||
return 2
|
||||
|
||||
if idx < nvmlDeviceGetCount():
|
||||
handle = nvmlDeviceGetHandleByIndex(idx)
|
||||
memInfo = nvmlDeviceGetMemoryInfo( handle )
|
||||
result = memInfo.total / (1024*1024*1024)
|
||||
nvmlShutdown()
|
||||
result = round(result)
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
memInfo = nvmlDeviceGetMemoryInfo( nvmlDeviceGetHandleByIndex(idx) )
|
||||
return round ( memInfo.total / (1024*1024*1024) )
|
||||
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def getBestDeviceIdx():
|
||||
if not hasNVML:
|
||||
return 0
|
||||
|
||||
idx = -1
|
||||
try:
|
||||
nvmlInit()
|
||||
idx_mem = 0
|
||||
for i in range(0, nvmlDeviceGetCount() ):
|
||||
handle = nvmlDeviceGetHandleByIndex(i)
|
||||
memInfo = nvmlDeviceGetMemoryInfo( handle )
|
||||
for i in range( nvmlDeviceGetCount() ):
|
||||
memInfo = nvmlDeviceGetMemoryInfo( nvmlDeviceGetHandleByIndex(i) )
|
||||
if memInfo.total > idx_mem:
|
||||
idx = i
|
||||
idx_mem = memInfo.total
|
||||
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return idx
|
||||
|
||||
@staticmethod
|
||||
def getWorstDeviceIdx():
|
||||
idx = -1
|
||||
try:
|
||||
nvmlInit()
|
||||
if not hasNVML:
|
||||
return 0
|
||||
|
||||
idx = -1
|
||||
idx_mem = sys.maxsize
|
||||
for i in range(0, nvmlDeviceGetCount() ):
|
||||
handle = nvmlDeviceGetHandleByIndex(i)
|
||||
memInfo = nvmlDeviceGetMemoryInfo( handle )
|
||||
for i in range( nvmlDeviceGetCount() ):
|
||||
memInfo = nvmlDeviceGetMemoryInfo( nvmlDeviceGetHandleByIndex(i) )
|
||||
if memInfo.total < idx_mem:
|
||||
idx = i
|
||||
idx_mem = memInfo.total
|
||||
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return idx
|
||||
|
||||
@staticmethod
|
||||
def isValidDeviceIdx(idx):
|
||||
result = False
|
||||
try:
|
||||
nvmlInit()
|
||||
result = (idx < nvmlDeviceGetCount())
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
if not hasNVML:
|
||||
return (idx == 0)
|
||||
|
||||
return (idx < nvmlDeviceGetCount())
|
||||
|
||||
@staticmethod
|
||||
def getDeviceIdxsEqualModel(idx):
|
||||
result = []
|
||||
try:
|
||||
nvmlInit()
|
||||
idx_name = nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode()
|
||||
if not hasNVML:
|
||||
return [0] if idx == 0 else []
|
||||
|
||||
for i in range(0, nvmlDeviceGetCount() ):
|
||||
result = []
|
||||
idx_name = nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode()
|
||||
for i in range( nvmlDeviceGetCount() ):
|
||||
if nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)).decode() == idx_name:
|
||||
result.append (i)
|
||||
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def getDeviceName (idx):
|
||||
result = 'Generic GeForce GPU'
|
||||
try:
|
||||
nvmlInit()
|
||||
if not hasNVML:
|
||||
return 'Generic GeForce GPU'
|
||||
|
||||
if idx < nvmlDeviceGetCount():
|
||||
result = nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode()
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result
|
||||
return nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode()
|
||||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def getDeviceComputeCapability(idx):
|
||||
if not hasNVML:
|
||||
return 99 if idx == 0 else 0
|
||||
|
||||
result = 0
|
||||
try:
|
||||
nvmlInit()
|
||||
if idx < nvmlDeviceGetCount():
|
||||
result = nvmlDeviceGetCudaComputeCapability(nvmlDeviceGetHandleByIndex(idx))
|
||||
nvmlShutdown()
|
||||
except:
|
||||
pass
|
||||
return result[0] * 10 + result[1]
|
||||
Loading…
Add table
Add a link
Reference in a new issue