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

fix xseg training

Browse source
This commit is contained in:
iperov 2020-12-11 15:47:11 +04:00
parent 35945b257c
commit 254a7cf5cf
2 changed files with 62 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

20
facelib/XSegNet.py
View file

@ -30,6 +30,9 @@ class XSegNet(object):
nn.initialize(data_format=data_format)
tf = nn.tf
model_name = f'{name}_{resolution}'
self.model_filename_list = []
with tf.device ('/CPU:0'):
#Place holders on CPU
self.input_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,3) )
@ -39,18 +42,17 @@ class XSegNet(object):
with tf.device ('/CPU:0' if place_model_on_cpu else '/GPU:0'):
self.model = nn.XSeg(3, 32, 1, name=name)
self.model_weights = self.model.get_weights()
if training:
if optimizer is None:
raise ValueError("Optimizer should be provided for training mode.")
self.opt = optimizer
self.opt.initialize_variables (self.model_weights, vars_on_cpu=place_model_on_cpu)
self.model_filename_list += [ [self.opt, f'{model_name}_opt.npy' ] ]
model_name = f'{name}_{resolution}'
self.model_filename_list = [ [self.model, f'{model_name}.npy'] ]
if training:
if optimizer is None:
raise ValueError("Optimizer should be provided for training mode.")
self.model_filename_list += [ [self.model, f'{model_name}.npy'] ]
self.opt = optimizer
self.opt.initialize_variables (self.model_weights, vars_on_cpu=place_model_on_cpu)
self.model_filename_list += [ [self.opt, f'{model_name}_opt.npy' ] ]
else:
if not training:
with tf.device ('/CPU:0' if run_on_cpu else '/GPU:0'):
_, pred = self.model(self.input_t)

16
models/Model_XSeg/Model.py
View file

@ -81,8 +81,9 @@ class XSegModel(ModelBase):
gpu_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
@ -100,10 +101,10 @@ class XSegModel(ModelBase):
# Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device):
pred = nn.concat(gpu_pred_list, 0)
loss = tf.reduce_mean(gpu_losses)
with tf.device(f'/CPU:0'):
#with tf.device (models_opt_device):
pred = tf.concat(gpu_pred_list, 0)
loss = tf.concat(gpu_losses, 0)
loss_gv_op = self.model.opt.get_update_op (nn.average_gv_list (gpu_loss_gvs))
@ -157,12 +158,9 @@ class XSegModel(ModelBase):
#override
def onTrainOneIter(self):
image_np, mask_np = self.generate_next_samples()[0]
loss = self.train (image_np, mask_np)
return ( ('loss', loss ), )
return ( ('loss', np.mean(loss) ), )
#override
def onGetPreview(self, samples):