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.should_save_preview_history():
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:
previews = self.get_static_previews()
for i in range(len(previews)):
name, bgr = previews[i]
path = self.preview_history_path / name
plist += [ ( bgr, str ( path / ( f'{self.iter:07d}.jpg') ) ) ]
if not io.is_colab():
plist += [ ( bgr, str ( path / ( '_last.jpg' ) )) ]
if len(plist) != 0:
self.get_preview_history_writer().post(plist, self.loss_history, self.iter)
self.iter += 1
return self.iter, iter_time
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 (not io.is_colab() and self.iter % 10 == 0) or \
(io.is_colab() and self.iter % 1 == 0):
print('doing preview')
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)
print(f'writting preview to {filepath}')
cv2_imwrite (filepath, img )
self.iter += 1
return self.iter, iter_time
def ask_write_preview_history(self, default_value=False):
default_write_preview_history = self.load_or_def_option('write_preview_history', default_value)
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 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 (not io.is_colab() and self.iter % 10 == 0) or \
(io.is_colab() and self.iter % 100 == 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:
previews = self.get_static_previews()
for i in range(len(previews)):
name, bgr = previews[i]
path = self.preview_history_path / name
path.mkdir(parents=True, exist_ok=True)
plist += [(bgr, str(path / (f'{self.iter:07d}.jpg')))]
if not io.is_colab():
plist += [(bgr, str(path / ('_last.jpg')))]
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 should_save_preview_history(self):
return (not io.is_colab() and self.iter % ( 10*(max(1,self.resolution // 64)) ) == 0) or \
(io.is_colab() and self.iter % 100 == 0)
def main(model_class_name=None,
saved_models_path=None,
training_data_src_path=None,
force_model_name=None,
input_path=None,
output_path=None,
output_mask_path=None,
aligned_path=None,
force_gpu_idxs=None,
cpu_only=None):
io.log_info("Running merger.\r\n")
try:
if not input_path.exists():
io.log_err('Input directory not found. Please ensure it exists.')
return
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
if not output_mask_path.exists():
output_mask_path.mkdir(parents=True, exist_ok=True)
if not saved_models_path.exists():
io.log_err('Model directory not found. Please ensure it exists.')
return
# Initialize model
import models
model = models.import_model(model_class_name)(
is_training=False,
saved_models_path=saved_models_path,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only)
predictor_func, predictor_input_shape, cfg = model.get_MergerConfig()
# Preparing MP functions
predictor_func = MPFunc(predictor_func)
run_on_cpu = len(nn.getCurrentDeviceConfig().devices) == 0
xseg_256_extract_func = MPClassFuncOnDemand(
XSegNet,
'extract',
name='XSeg',
resolution=256,
weights_file_root=saved_models_path,
place_model_on_cpu=True,
run_on_cpu=run_on_cpu)
face_enhancer_func = MPClassFuncOnDemand(FaceEnhancer,
'enhance',
place_model_on_cpu=True,
run_on_cpu=run_on_cpu)
is_interactive = io.input_bool("Use interactive merger?",
True) if not io.is_colab() else False
# if not is_interactive:
# cfg.ask_settings()
subprocess_count = multiprocessing.cpu_count()
# subprocess_count = io.input_int("Number of workers?", max(8, multiprocessing.cpu_count()),
# valid_range=[1, multiprocessing.cpu_count()], help_message="Specify the number of threads to process. A low value may affect performance. A high value may result in memory error. The value may not be greater than CPU cores." )
input_path_image_paths = pathex.get_image_paths(input_path)
if cfg.type == MergerConfig.TYPE_MASKED:
if not aligned_path.exists():
io.log_err(
'Aligned directory not found. Please ensure it exists.')
return
packed_samples = None
try:
packed_samples = samplelib.PackedFaceset.load(aligned_path)
except:
io.log_err(
f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}"
)
if packed_samples is not None:
io.log_info("Using packed faceset.")
def generator():
for sample in io.progress_bar_generator(
packed_samples, "Collecting alignments"):
filepath = Path(sample.filename)
yield filepath, DFLIMG.load(
filepath,
loader_func=lambda x: sample.read_raw_file())
else:
def generator():
for filepath in io.progress_bar_generator(
pathex.get_image_paths(aligned_path),
"Collecting alignments"):
filepath = Path(filepath)
yield filepath, DFLIMG.load(filepath)
alignments = {}
multiple_faces_detected = False
for filepath, dflimg in generator():
if dflimg is None or not dflimg.has_data():
io.log_err(f"{filepath.name} is not a dfl image file")
continue
source_filename = dflimg.get_source_filename()
if source_filename is None:
continue
source_filepath = Path(source_filename)
source_filename_stem = source_filepath.stem
if source_filename_stem not in alignments.keys():
alignments[source_filename_stem] = []
alignments_ar = alignments[source_filename_stem]
alignments_ar.append(
(dflimg.get_source_landmarks(), filepath, source_filepath))
if len(alignments_ar) > 1:
multiple_faces_detected = True
if multiple_faces_detected:
io.log_info("")
io.log_info(
"Warning: multiple faces detected. Only one alignment file should refer one source file."
)
io.log_info("")
for a_key in list(alignments.keys()):
a_ar = alignments[a_key]
if len(a_ar) > 1:
for _, filepath, source_filepath in a_ar:
io.log_info(
f"alignment {filepath.name} refers to {source_filepath.name} "
)
io.log_info("")
alignments[a_key] = [a[0] for a in a_ar]
if multiple_faces_detected:
io.log_info(
"It is strongly recommended to process the faces separatelly."
)
io.log_info(
"Use 'recover original filename' to determine the exact duplicates."
)
io.log_info("")
frames = [
InteractiveMergerSubprocessor.Frame(frame_info=FrameInfo(
filepath=Path(p),
landmarks_list=alignments.get(Path(p).stem, None)))
for p in input_path_image_paths
]
if multiple_faces_detected:
io.log_info(
"Warning: multiple faces detected. Motion blur will not be used."
)
io.log_info("")
else:
s = 256
local_pts = [(s // 2 - 1, s // 2 - 1),
(s // 2 - 1, 0)] #center+up
frames_len = len(frames)
for i in io.progress_bar_generator(range(len(frames)),
"Computing motion vectors"):
fi_prev = frames[max(0, i - 1)].frame_info
fi = frames[i].frame_info
fi_next = frames[min(i + 1, frames_len - 1)].frame_info
if len(fi_prev.landmarks_list) == 0 or \
len(fi.landmarks_list) == 0 or \
len(fi_next.landmarks_list) == 0:
continue
mat_prev = LandmarksProcessor.get_transform_mat(
fi_prev.landmarks_list[0], s, face_type=FaceType.FULL)
mat = LandmarksProcessor.get_transform_mat(
fi.landmarks_list[0], s, face_type=FaceType.FULL)
mat_next = LandmarksProcessor.get_transform_mat(
fi_next.landmarks_list[0], s, face_type=FaceType.FULL)
pts_prev = LandmarksProcessor.transform_points(
local_pts, mat_prev, True)
pts = LandmarksProcessor.transform_points(
local_pts, mat, True)
pts_next = LandmarksProcessor.transform_points(
local_pts, mat_next, True)
prev_vector = pts[0] - pts_prev[0]
next_vector = pts_next[0] - pts[0]
motion_vector = pts_next[0] - pts_prev[0]
fi.motion_power = npla.norm(motion_vector)
motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array(
[0, 0], dtype=np.float32)
fi.motion_deg = -math.atan2(
motion_vector[1], motion_vector[0]) * 180 / math.pi
if len(frames) == 0:
io.log_info("No frames to merge in input_dir.")
else:
if False:
pass
else:
InteractiveMergerSubprocessor(
is_interactive=is_interactive,
merger_session_filepath=model.get_strpath_storage_for_file(
'merger_session.dat'),
predictor_func=predictor_func,
predictor_input_shape=predictor_input_shape,
face_enhancer_func=face_enhancer_func,
xseg_256_extract_func=xseg_256_extract_func,
merger_config=cfg,
frames=frames,
frames_root_path=input_path,
output_path=output_path,
output_mask_path=output_mask_path,
model_iter=model.get_iter(),
subprocess_count=subprocess_count,
).run()
model.finalize()
except Exception as e:
print(traceback.format_exc())
def trainerThread(s2c,
c2s,
e,
model_class_name=None,
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=None,
execute_programs=None,
debug=False,
**kwargs):
while True:
try:
start_time = time.time()
save_interval_min = 15
if not training_data_src_path.exists():
io.log_err('Training data src directory does not exist.')
break
if not training_data_dst_path.exists():
io.log_err('Training data dst directory does not exist.')
break
if not saved_models_path.exists():
saved_models_path.mkdir(exist_ok=True)
model = models.import_model(model_class_name)(
is_training=True,
saved_models_path=saved_models_path,
training_data_src_path=training_data_src_path,
training_data_dst_path=training_data_dst_path,
pretraining_data_path=pretraining_data_path,
pretrained_model_path=pretrained_model_path,
no_preview=no_preview,
force_model_name=force_model_name,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only,
debug=debug,
)
is_reached_goal = model.is_reached_iter_goal()
shared_state = {'after_save': False}
loss_string = ""
save_iter = model.get_iter()
def model_save():
if not debug and not is_reached_goal:
io.log_info("Saving....", end='\r')
model.save()
shared_state['after_save'] = True
def send_preview():
if not debug:
previews = model.get_previews()
c2s.put({
'op': 'show',
'previews': previews,
'iter': model.get_iter(),
'loss_history': model.get_loss_history().copy()
})
else:
previews = [('debug, press update for new',
model.debug_one_iter())]
c2s.put({'op': 'show', 'previews': previews})
e.set() #Set the GUI Thread as Ready
if model.get_target_iter() != 0:
if is_reached_goal:
io.log_info(
'Model already trained to target iteration. You can use preview.'
)
else:
io.log_info(
'Starting. Target iteration: %d. Press "Enter" to stop training and save model.'
% (model.get_target_iter()))
else:
io.log_info(
'Starting. Press "Enter" to stop training and save model.')
last_save_time = time.time()
execute_programs = [[x[0], x[1], time.time()]
for x in execute_programs]
for i in itertools.count(0, 1):
if not debug:
cur_time = time.time()
for x in execute_programs:
prog_time, prog, last_time = x
exec_prog = False
if prog_time > 0 and (cur_time -
start_time) >= prog_time:
x[0] = 0
exec_prog = True
elif prog_time < 0 and (cur_time -
last_time) >= -prog_time:
x[2] = cur_time
exec_prog = True
if exec_prog:
try:
exec(prog)
except Exception as e:
print("Unable to execute program: %s" % (prog))
if not is_reached_goal:
if model.get_iter() == 0:
io.log_info("")
io.log_info(
"Trying to do the first iteration. If an error occurs, reduce the model parameters."
)
io.log_info("")
iter, iter_time = model.train_one_iter()
loss_history = model.get_loss_history()
time_str = time.strftime("[%H:%M:%S]")
if iter_time >= 10:
loss_string = "{0}[#{1:06d}][{2:.5s}s]".format(
time_str, iter, '{:0.4f}'.format(iter_time))
else:
loss_string = "{0}[#{1:06d}][{2:04d}ms]".format(
time_str, iter, int(iter_time * 1000))
if shared_state['after_save']:
shared_state['after_save'] = False
last_save_time = time.time()
mean_loss = np.mean([
np.array(loss_history[i])
for i in range(save_iter, iter)
],
axis=0)
for loss_value in mean_loss:
loss_string += "[%.4f]" % (loss_value)
io.log_info(loss_string)
save_iter = iter
else:
for loss_value in loss_history[-1]:
loss_string += "[%.4f]" % (loss_value)
if io.is_colab():
io.log_info('\r' + loss_string, end='')
else:
io.log_info(loss_string, end='\r')
if model.get_iter() == 1:
model_save()
if model.get_target_iter(
) != 0 and model.is_reached_iter_goal():
io.log_info('Reached target iteration.')
model_save()
is_reached_goal = True
io.log_info('You can use preview now.')
if not is_reached_goal and (time.time() - last_save_time
) >= save_interval_min * 60:
model_save()
send_preview()
if i == 0:
if is_reached_goal:
model.pass_one_iter()
send_preview()
if debug:
time.sleep(0.005)
while not s2c.empty():
input = s2c.get()
op = input['op']
if op == 'save':
model_save()
elif op == 'preview':
if is_reached_goal:
model.pass_one_iter()
send_preview()
elif op == 'close':
model_save()
i = -1
break
if i == -1:
break
model.finalize()
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
break
c2s.put({'op': 'close'})
def main(model_class_name=None,
saved_models_path=None,
training_data_src_path=None,
force_model_name=None,
input_path=None,
output_path=None,
aligned_path=None,
force_gpu_idxs=None,
cpu_only=None):
io.log_info("Running merger.\r\n")
try:
if not input_path.exists():
io.log_err('Input directory not found. Please ensure it exists.')
return
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
if not saved_models_path.exists():
io.log_err('Model directory not found. Please ensure it exists.')
return
is_interactive = io.input_bool("Use interactive merger?",
True) if not io.is_colab() else False
import models
model = models.import_model(model_class_name)(
is_training=False,
saved_models_path=saved_models_path,
training_data_src_path=training_data_src_path,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only)
merger_session_filepath = model.get_strpath_storage_for_file(
'merger_session.dat')
predictor_func, predictor_input_shape, cfg = model.get_MergerConfig()
if not is_interactive:
cfg.ask_settings()
input_path_image_paths = pathex.get_image_paths(input_path)
if cfg.type == MergerConfig.TYPE_MASKED:
if not aligned_path.exists():
io.log_err(
'Aligned directory not found. Please ensure it exists.')
return
packed_samples = None
try:
packed_samples = samplelib.PackedFaceset.load(aligned_path)
except:
io.log_err(
f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}"
)
if packed_samples is not None:
io.log_info("Using packed faceset.")
def generator():
for sample in io.progress_bar_generator(
packed_samples, "Collecting alignments"):
filepath = Path(sample.filename)
yield DFLIMG.load(
filepath,
loader_func=lambda x: sample.read_raw_file())
else:
def generator():
for filepath in io.progress_bar_generator(
pathex.get_image_paths(aligned_path),
"Collecting alignments"):
filepath = Path(filepath)
yield DFLIMG.load(filepath)
alignments = {}
multiple_faces_detected = False
for dflimg in generator():
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
continue
source_filename = dflimg.get_source_filename()
if source_filename is None or source_filename == "_":
continue
source_filename = Path(source_filename)
source_filename_stem = source_filename.stem
if source_filename_stem not in alignments.keys():
alignments[source_filename_stem] = []
alignments_ar = alignments[source_filename_stem]
alignments_ar.append(dflimg.get_source_landmarks())
if len(alignments_ar) > 1:
multiple_faces_detected = True
if multiple_faces_detected:
io.log_info(
"Warning: multiple faces detected. Strongly recommended to process them separately."
)
frames = [
MergeSubprocessor.Frame(frame_info=FrameInfo(
filepath=Path(p),
landmarks_list=alignments.get(Path(p).stem, None)))
for p in input_path_image_paths
]
if multiple_faces_detected:
io.log_info(
"Warning: multiple faces detected. Motion blur will not be used."
)
else:
s = 256
local_pts = [(s // 2 - 1, s // 2 - 1),
(s // 2 - 1, 0)] #center+up
frames_len = len(frames)
for i in io.progress_bar_generator(range(len(frames)),
"Computing motion vectors"):
fi_prev = frames[max(0, i - 1)].frame_info
fi = frames[i].frame_info
fi_next = frames[min(i + 1, frames_len - 1)].frame_info
if len(fi_prev.landmarks_list) == 0 or \
len(fi.landmarks_list) == 0 or \
len(fi_next.landmarks_list) == 0:
continue
mat_prev = LandmarksProcessor.get_transform_mat(
fi_prev.landmarks_list[0], s, face_type=FaceType.FULL)
mat = LandmarksProcessor.get_transform_mat(
fi.landmarks_list[0], s, face_type=FaceType.FULL)
mat_next = LandmarksProcessor.get_transform_mat(
fi_next.landmarks_list[0], s, face_type=FaceType.FULL)
pts_prev = LandmarksProcessor.transform_points(
local_pts, mat_prev, True)
pts = LandmarksProcessor.transform_points(
local_pts, mat, True)
pts_next = LandmarksProcessor.transform_points(
local_pts, mat_next, True)
prev_vector = pts[0] - pts_prev[0]
next_vector = pts_next[0] - pts[0]
motion_vector = pts_next[0] - pts_prev[0]
fi.motion_power = npla.norm(motion_vector)
motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array(
[0, 0], dtype=np.float32)
fi.motion_deg = -math.atan2(
motion_vector[1], motion_vector[0]) * 180 / math.pi
elif cfg.type == MergerConfig.TYPE_FACE_AVATAR:
filesdata = []
for filepath in io.progress_bar_generator(input_path_image_paths,
"Collecting info"):
filepath = Path(filepath)
dflimg = DFLIMG.load(filepath)
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
continue
filesdata += [
(FrameInfo(filepath=filepath,
landmarks_list=[dflimg.get_landmarks()]),
dflimg.get_source_filename())
]
filesdata = sorted(
filesdata,
key=operator.itemgetter(1)) #sort by source_filename
frames = []
filesdata_len = len(filesdata)
for i in range(len(filesdata)):
frame_info = filesdata[i][0]
prev_temporal_frame_infos = []
next_temporal_frame_infos = []
for t in range(cfg.temporal_face_count):
prev_frame_info = filesdata[max(i - t, 0)][0]
next_frame_info = filesdata[min(i + t,
filesdata_len - 1)][0]
prev_temporal_frame_infos.insert(0, prev_frame_info)
next_temporal_frame_infos.append(next_frame_info)
frames.append(
MergeSubprocessor.Frame(
prev_temporal_frame_infos=prev_temporal_frame_infos,
frame_info=frame_info,
next_temporal_frame_infos=next_temporal_frame_infos))
if len(frames) == 0:
io.log_info("No frames to merge in input_dir.")
else:
MergeSubprocessor(is_interactive=is_interactive,
merger_session_filepath=merger_session_filepath,
predictor_func=predictor_func,
predictor_input_shape=predictor_input_shape,
merger_config=cfg,
frames=frames,
frames_root_path=input_path,
output_path=output_path,
model_iter=model.get_iter()).run()
model.finalize()
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
def trainerThread(s2c,
c2s,
e,
model_class_name=None,
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=None,
silent_start=False,
execute_programs=None,
debug=False,
**kwargs):
while True:
try:
start_time = time.time()
save_interval_min = 15
if not training_data_src_path.exists():
training_data_src_path.mkdir(exist_ok=True, parents=True)
if not training_data_dst_path.exists():
training_data_dst_path.mkdir(exist_ok=True, parents=True)
if not saved_models_path.exists():
saved_models_path.mkdir(exist_ok=True, parents=True)
model = models.import_model(model_class_name)(
is_training=True,
saved_models_path=saved_models_path,
training_data_src_path=training_data_src_path,
training_data_dst_path=training_data_dst_path,
pretraining_data_path=pretraining_data_path,
pretrained_model_path=pretrained_model_path,
no_preview=no_preview,
force_model_name=force_model_name,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only,
silent_start=silent_start,
debug=debug,
)
is_reached_goal = model.is_reached_iter_goal()
shared_state = {'after_save': False}
loss_string = ""
save_iter = model.get_iter()
def model_save():
if not debug and not is_reached_goal:
io.log_info("保存....", end='\r')
model.save()
shared_state['after_save'] = True
def model_backup():
if not debug and not is_reached_goal:
model.create_backup()
def send_preview():
if not debug:
previews = model.get_previews()
c2s.put({
'op': 'show',
'previews': previews,
'iter': model.get_iter(),
'loss_history': model.get_loss_history().copy()
})
else:
previews = [('debug, press update for new',
model.debug_one_iter())]
c2s.put({'op': 'show', 'previews': previews})
e.set() #Set the GUI Thread as Ready
if model.get_target_iter() != 0:
if is_reached_goal:
io.log_info('模型训练次数已经达到设定值. 你可以通过窗口预览效果.')
else:
io.log_info('启动中. 目标迭代: %d. 按 "Enter" 停止训练并保存进度.' %
(model.get_target_iter()))
else:
io.log_info(
"\n 各种高性价比、参数模型\n 联系QQ：395267954\n=============================================\n"
)
io.log_info(
'启动中. 按 "Enter" 停止训练并保存进度。\n\n保存时间|迭代次数|单次时间|源损失|目标损失')
last_save_time = time.time()
execute_programs = [[x[0], x[1], time.time()]
for x in execute_programs]
for i in itertools.count(0, 1):
if not debug:
cur_time = time.time()
for x in execute_programs:
prog_time, prog, last_time = x
exec_prog = False
if prog_time > 0 and (cur_time -
start_time) >= prog_time:
x[0] = 0
exec_prog = True
elif prog_time < 0 and (cur_time -
last_time) >= -prog_time:
x[2] = cur_time
exec_prog = True
if exec_prog:
try:
exec(prog)
except Exception as e:
print("Unable to execute program: %s" % (prog))
if not is_reached_goal:
if model.get_iter() == 0:
io.log_info("")
io.log_info("正在尝试运行第一个迭代. 如果出现错误, 请降低参数配置")
io.log_info("")
iter, iter_time = model.train_one_iter()
loss_history = model.get_loss_history()
time_str = time.strftime("[%H:%M:%S]")
if iter_time >= 10:
loss_string = "{0}[#{1:06d}][{2:.5s}s]".format(
time_str, iter, '{:0.4f}'.format(iter_time))
else:
loss_string = "{0}[#{1:06d}][{2:04d}ms]".format(
time_str, iter, int(iter_time * 1000))
if shared_state['after_save']:
shared_state['after_save'] = False
mean_loss = np.mean(loss_history[save_iter:iter],
axis=0)
for loss_value in mean_loss:
loss_string += "[%.4f]" % (loss_value)
io.log_info(loss_string)
save_iter = iter
else:
for loss_value in loss_history[-1]:
loss_string += "[%.4f]" % (loss_value)
if io.is_colab():
io.log_info('\r' + loss_string, end='')
else:
io.log_info(loss_string, end='\r')
if model.get_iter() == 1:
model_save()
if model.get_target_iter(
) != 0 and model.is_reached_iter_goal():
io.log_info('达到目标迭代.')
model_save()
is_reached_goal = True
io.log_info('可以使用预览窗口.')
if not is_reached_goal and (time.time() - last_save_time
) >= save_interval_min * 60:
last_save_time += save_interval_min * 60
model_save()
send_preview()
if i == 0:
if is_reached_goal:
model.pass_one_iter()
send_preview()
if debug:
time.sleep(0.005)
while not s2c.empty():
input = s2c.get()
op = input['op']
if op == 'save':
model_save()
elif op == 'backup':
model_backup()
elif op == 'preview':
if is_reached_goal:
model.pass_one_iter()
send_preview()
elif op == 'close':
model_save()
i = -1
break
if i == -1:
break
model.finalize()
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
break
c2s.put({'op': 'close'})
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 __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,
force_model_class_name=None,
**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_class_name is None:
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_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", "new")
self.model_name = self.model_name.replace('_', ' ')
break
self.model_name = self.model_name + '_' + self.model_class_name
else:
self.model_name = force_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_hour = self.options.get('autobackup_hour', 0)
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.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')
self.update_sample_for_preview(choose_preview_history=self.choose_preview_history)
if self.autobackup_hour != 0:
self.autobackup_start_time = time.time()
if not self.autobackups_path.exists():
self.autobackups_path.mkdir(exist_ok=True)
io.log_info( self.get_summary_text() )
def trainerThread(s2c,
c2s,
e,
model_class_name=None,
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=None,
silent_start=False,
execute_programs=None,
debug=False,
target_iter=None,
precision=None,
bs_per_gpu=None,
use_amp=None,
opt=None,
lr=None,
decay_step=None,
config_file=None,
**kwargs):
while True:
try:
start_time = time.time()
save_interval_min = 15
if not training_data_src_path.exists():
training_data_src_path.mkdir(exist_ok=True, parents=True)
if not training_data_dst_path.exists():
training_data_dst_path.mkdir(exist_ok=True, parents=True)
if not saved_models_path.exists():
saved_models_path.mkdir(exist_ok=True, parents=True)
model = models.import_model_tf1(model_class_name)(
is_training=True,
saved_models_path=saved_models_path,
training_data_src_path=training_data_src_path,
training_data_dst_path=training_data_dst_path,
pretraining_data_path=pretraining_data_path,
pretrained_model_path=pretrained_model_path,
no_preview=no_preview,
force_model_name=force_model_name,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only,
silent_start=silent_start,
debug=debug,
target_iter=target_iter,
precision=precision,
bs_per_gpu=bs_per_gpu,
use_amp=use_amp,
opt=opt,
lr=lr,
decay_step=decay_step,
config_file=config_file)
is_reached_goal = model.is_reached_iter_goal()
shared_state = {'after_save': False}
loss_string = ""
save_iter = model.get_iter()
def model_save():
if not debug and not is_reached_goal:
io.log_info("Saving....", end='\r')
model.save()
shared_state['after_save'] = True
def model_backup():
if not debug and not is_reached_goal:
model.create_backup()
def send_preview():
if not debug:
previews = model.get_previews()
c2s.put({
'op': 'show',
'previews': previews,
'iter': model.get_iter(),
'loss_history': model.get_loss_history().copy()
})
else:
previews = [('debug, press update for new',
model.debug_one_iter())]
c2s.put({'op': 'show', 'previews': previews})
e.set() #Set the GUI Thread as Ready
if model.get_target_iter() != 0:
if is_reached_goal:
io.log_info(
'Model already trained to target iteration. You can use preview.'
)
else:
io.log_info(
'Starting. Target iteration: %d. Press "Enter" to stop training and save model.'
% (model.get_target_iter()))
else:
io.log_info(
'Starting. Press "Enter" to stop training and save model.')
last_save_time = time.time()
execute_programs = [[x[0], x[1], time.time()]
for x in execute_programs]
tf = nn.tf
list_globals_vars = tf.get_collection(
tf.GraphKeys.GLOBAL_VARIABLES)
print('initializing variables ... ')
list_init = []
for x in [n for n in list_globals_vars]:
if 'Adam' in x.name or 'RMSProp' in x.name or 'beta' in x.name or 'loss_scale' in x.name or 'good_steps' in x.name:
if not nn.tf_sess.run(nn.tf.is_variable_initialized(x)):
list_init.append(x)
nn.tf_sess.run(tf.variables_initializer(list_init))
nn.tf_sess.run(model.global_step.initializer)
run_opts = tf.RunOptions(report_tensor_allocations_upon_oom=False)
print('done ')
# ( (warped_src, target_src, target_srcm_all), \
# (warped_dst, target_dst, target_dstm_all) ) = model.generate_next_samples()
for i in itertools.count(0, 1):
if not debug:
cur_time = time.time()
for x in execute_programs:
prog_time, prog, last_time = x
exec_prog = False
if prog_time > 0 and (cur_time -
start_time) >= prog_time:
x[0] = 0
exec_prog = True
elif prog_time < 0 and (cur_time -
last_time) >= -prog_time:
x[2] = cur_time
exec_prog = True
if exec_prog:
try:
exec(prog)
except Exception as e:
print("Unable to execute program: %s" % (prog))
if not is_reached_goal:
if model.get_iter() == 0:
io.log_info("")
io.log_info(
"Trying to do the first iteration. If an error occurs, reduce the model parameters."
)
io.log_info("")
t_start = time.time()
( (warped_src, target_src, target_srcm_all), \
(warped_dst, target_dst, target_dstm_all) ) = model.generate_next_samples()
list_loss = []
# Train auto-encoder
# Train different parts of the network in sequence
# More accurate gradient, Slower
if True:
_, src_loss, dst_loss, learning_rate = nn.tf_sess.run(
[
model.G_train_op, model.src_loss,
model.dst_loss, model.learning_rate
],
feed_dict={
model.warped_src: warped_src,
model.target_src: target_src,
model.target_srcm_all: target_srcm_all,
model.warped_dst: warped_dst,
model.target_dst: target_dst,
model.target_dstm_all: target_dstm_all
})
# _, src_loss, dst_loss, learning_rate, \
# gpu_src_code, gpu_src_inter_AB_code, gpu_dst_inter_B_code, \
# gpu_pred_src_src, gpu_pred_src_srcm = nn.tf_sess.run(
# [model.G_train_op, model.src_loss, model.dst_loss, model.learning_rate,
# model.gpu_src_code, model.gpu_src_inter_AB_code, model.gpu_dst_inter_B_code,
# model.gpu_pred_src_src, model.gpu_pred_src_srcm], feed_dict={
# model.warped_src :warped_src,
# model.target_src :target_src,
# model.target_srcm_all:target_srcm_all,
# model.warped_dst :warped_dst,
# model.target_dst :target_dst,
# model.target_dstm_all:target_dstm_all})
# print(gpu_src_code.shape)
# print(gpu_src_inter_AB_code.shape)
# print(gpu_dst_inter_B_code.shape)
# print(gpu_pred_src_src.shape)
# print(gpu_pred_src_srcm.shape)
# print('-----------------------------------------------------')
# Train face style
if model.options['true_face_power'] != 0:
_, _D_code_loss = nn.tf_sess.run(
[model.D_code_train_op, model.D_code_loss],
feed_dict={
model.warped_src: warped_src,
model.target_src: target_src,
model.target_srcm_all: target_srcm_all,
model.warped_dst: warped_dst,
model.target_dst: target_dst,
model.target_dstm_all: target_dstm_all
})
# Train GAN
if model.options['gan_power'] != 0:
_, D_src_dst_loss = nn.tf_sess.run(
[
model.D_src_dst_train_op,
model.D_src_dst_loss
],
feed_dict={
model.warped_src: warped_src,
model.target_src: target_src,
model.target_srcm_all: target_srcm_all,
model.warped_dst: warped_dst,
model.target_dst: target_dst,
model.target_dstm_all: target_dstm_all
})
# Train different parts of the network in parallel
# Less accurate gradient, Faster
if False:
if model.options[
'gan_power'] == 0 and model.options[
'true_face_power'] == 0:
_G_train_op, src_loss, dst_loss = nn.tf_sess.run(
[
model.G_train_op, model.src_loss,
model.dst_loss
],
feed_dict={
model.warped_src: warped_src,
model.target_src: target_src,
model.target_srcm_all: target_srcm_all,
model.warped_dst: warped_dst,
model.target_dst: target_dst,
model.target_dstm_all: target_dstm_all
})
elif model.options['gan_power'] != 0:
_G_train_op, _D_src_dst_train_op, src_loss, dst_loss, D_src_dst_loss = nn.tf_sess.run(
[
model.G_train_op,
model.D_src_dst_train_op,
model.src_loss, model.dst_loss,
model.D_src_dst_loss
],
feed_dict={
model.warped_src: warped_src,
model.target_src: target_src,
model.target_srcm_all: target_srcm_all,
model.warped_dst: warped_dst,
model.target_dst: target_dst,
model.target_dstm_all: target_dstm_all
})
elif model.options['true_face_power'] != 0:
_G_train_op, _D_code_train_op, src_loss, dst_loss, D_code_loss = nn.tf_sess.run(
[
model.G_train_op,
model.D_code_train_op, model.src_loss,
model.dst_loss, model.D_code_loss
],
feed_dict={
model.warped_src: warped_src,
model.target_src: target_src,
model.target_srcm_all: target_srcm_all,
model.warped_dst: warped_dst,
model.target_dst: target_dst,
model.target_dstm_all: target_dstm_all
})
src_loss = 2.0 if math.isnan(src_loss) else src_loss
dst_loss = 2.0 if math.isnan(dst_loss) else dst_loss
list_loss = [float(src_loss), float(dst_loss)]
model.loss_history.append(list_loss)
model.iter += 1
iter_time = time.time() - t_start
iter = model.get_iter()
if (not io.is_colab() and iter % 200 == 0) or \
(io.is_colab() and iter % 200 == 0):
plist = []
if io.is_colab():
previews = model.get_previews()
for i in range(len(previews)):
name, bgr = previews[i]
plist += [
(bgr,
model.get_strpath_storage_for_file(
'preview_%s.jpg' % (name)))
]
if model.write_preview_history:
previews = model.get_static_previews()
for i in range(len(previews)):
name, bgr = previews[i]
path = model.preview_history_path / name
path.mkdir(parents=True, exist_ok=True)
plist += [
(bgr, str(path / (f'{iter:07d}.jpg')))
]
if not io.is_colab():
plist += [(bgr,
str(path / ('_last.jpg')))]
for preview, filepath in plist:
preview_lh = model.get_loss_history_preview(
model.loss_history, iter, preview.shape[1],
preview.shape[2])
img = (np.concatenate([preview_lh, preview],
axis=0) * 255).astype(
np.uint8)
cv2_imwrite(filepath, img)
# iter, iter_time = model.train_one_iter()
loss_history = model.get_loss_history()
time_str = time.strftime("[%H:%M:%S]")
if iter_time >= 10:
loss_string = "{0}[#{1:06d}][{2:.5s}s]".format(
time_str, iter, '{:0.4f}'.format(iter_time))
else:
loss_string = "{0}[#{1:06d}][{2:04d}ms]".format(
time_str, iter, int(iter_time * 1000))
if shared_state['after_save']:
shared_state['after_save'] = False
mean_loss = np.mean(loss_history[save_iter:iter],
axis=0)
for loss_value in mean_loss:
loss_string += "[%.5f]" % (loss_value)
io.log_info(loss_string)
save_iter = iter
else:
for loss_value in loss_history[-1]:
loss_string += "[%.5f]" % (loss_value)
if io.is_colab():
io.log_info('\r' + loss_string, end='')
else:
io.log_info(loss_string, end='\r')
if model.get_iter() == 1:
model_save()
if model.get_target_iter(
) != 0 and model.is_reached_iter_goal():
io.log_info('\nReached target iteration.')
model_save()
is_reached_goal = True
break
io.log_info('You can use preview now.')
if not is_reached_goal and (time.time() - last_save_time
) >= save_interval_min * 60:
last_save_time += save_interval_min * 60
model_save()
send_preview()
if i == 0:
if is_reached_goal:
model.pass_one_iter()
send_preview()
if debug:
time.sleep(0.005)
while not s2c.empty():
input = s2c.get()
op = input['op']
if op == 'save':
model_save()
elif op == 'backup':
model_backup()
elif op == 'preview':
if is_reached_goal:
model.pass_one_iter()
send_preview()
elif op == 'close':
model_save()
i = -1
break
if i == -1:
break
model.finalize()
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
break
c2s.put({'op': 'close'})
def should_save_preview_history(self):
return (not io.is_colab() and self.iter % 10 == 0) or (io.is_colab() and self.iter % 100 == 0)
def trainerThread(s2c,
c2s,
e,
model_class_name=None,
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=None,
silent_start=False,
execute_programs=None,
debug=False,
tensorboard_dir=None,
start_tensorboard=False,
**kwargs):
while True:
try:
start_time = time.time()
save_interval_min = 15
tensorboard_preview_interval_min = 5
if not training_data_src_path.exists():
training_data_src_path.mkdir(exist_ok=True, parents=True)
if not training_data_dst_path.exists():
training_data_dst_path.mkdir(exist_ok=True, parents=True)
if not saved_models_path.exists():
saved_models_path.mkdir(exist_ok=True, parents=True)
model = models.import_model(model_class_name)(
is_training=True,
saved_models_path=saved_models_path,
training_data_src_path=training_data_src_path,
training_data_dst_path=training_data_dst_path,
pretraining_data_path=pretraining_data_path,
pretrained_model_path=pretrained_model_path,
no_preview=no_preview,
force_model_name=force_model_name,
force_gpu_idxs=force_gpu_idxs,
cpu_only=cpu_only,
silent_start=silent_start,
debug=debug,
)
is_reached_goal = model.is_reached_iter_goal()
if tensorboard_dir is not None:
c2s.put({
'op': 'tb',
'action': 'init',
'model_name': model.model_name,
'tensorboard_dir': tensorboard_dir,
'start_tensorboard': start_tensorboard
})
shared_state = {'after_save': False}
loss_string = ""
save_iter = model.get_iter()
def model_save():
if not debug and not is_reached_goal:
io.log_info("Saving....", end='\r')
model.save()
shared_state['after_save'] = True
def model_backup():
if not debug and not is_reached_goal:
model.create_backup()
def log_step(step, step_time, src_loss, dst_loss):
c2s.put({
'op': 'tb',
'action': 'step',
'step': step,
'step_time': step_time,
'src_loss': src_loss,
'dst_loss': dst_loss
})
def log_previews(step, previews, static_previews):
c2s.put({
'op': 'tb',
'action': 'preview',
'step': step,
'previews': previews,
'static_previews': static_previews
})
def send_preview():
if not debug:
previews = model.get_previews()
c2s.put({
'op': 'show',
'previews': previews,
'iter': model.get_iter(),
'loss_history': model.get_loss_history().copy()
})
else:
previews = [('debug, press update for new',
model.debug_one_iter())]
c2s.put({'op': 'show', 'previews': previews})
e.set() #Set the GUI Thread as Ready
if model.get_target_iter() != 0:
if is_reached_goal:
io.log_info(
'Model already trained to target iteration. You can use preview.'
)
else:
io.log_info(
'Starting. Target iteration: %d. Press "Enter" to stop training and save model.'
% (model.get_target_iter()))
else:
io.log_info(
'Starting. Press "Enter" to stop training and save model.')
last_save_time = time.time()
last_preview_time = time.time()
execute_programs = [[x[0], x[1], time.time()]
for x in execute_programs]
for i in itertools.count(0, 1):
if not debug:
cur_time = time.time()
for x in execute_programs:
prog_time, prog, last_time = x
exec_prog = False
if prog_time > 0 and (cur_time -
start_time) >= prog_time:
x[0] = 0
exec_prog = True
elif prog_time < 0 and (cur_time -
last_time) >= -prog_time:
x[2] = cur_time
exec_prog = True
if exec_prog:
try:
exec(prog)
except Exception as e:
print("Unable to execute program: %s" % (prog))
if not is_reached_goal:
if model.get_iter() == 0:
io.log_info("")
io.log_info(
"Trying to do the first iteration. If an error occurs, reduce the model parameters."
)
io.log_info("")
if sys.platform[0:3] == 'win':
io.log_info("!!!")
io.log_info(
"Windows 10 users IMPORTANT notice. You should set this setting in order to work correctly."
)
io.log_info("https://i.imgur.com/B7cmDCB.jpg")
io.log_info("!!!")
iter, iter_time = model.train_one_iter()
loss_history = model.get_loss_history()
time_str = time.strftime("[%H:%M:%S]")
if iter_time >= 10:
loss_string = "{0}[#{1:06d}][{2:.5s}s]".format(
time_str, iter, '{:0.4f}'.format(iter_time))
else:
loss_string = "{0}[#{1:06d}][{2:04d}ms]".format(
time_str, iter, int(iter_time * 1000))
if shared_state['after_save']:
shared_state['after_save'] = False
mean_loss = np.mean(loss_history[save_iter:iter],
axis=0)
for loss_value in mean_loss:
loss_string += "[%.4f]" % (loss_value)
io.log_info(loss_string)
save_iter = iter
else:
for loss_value in loss_history[-1]:
loss_string += "[%.4f]" % (loss_value)
if io.is_colab():
io.log_info('\r' + loss_string, end='')
else:
io.log_info(loss_string, end='\r')
loss_entry = loss_history[-1]
log_step(
iter, iter_time, loss_entry[0],
loss_entry[1] if len(loss_entry) > 1 else None)
if model.get_iter() == 1:
model_save()
if model.get_target_iter(
) != 0 and model.is_reached_iter_goal():
io.log_info('Reached target iteration.')
model_save()
is_reached_goal = True
io.log_info('You can use preview now.')
if not is_reached_goal and (
time.time() - last_preview_time
) >= tensorboard_preview_interval_min * 60:
last_preview_time += tensorboard_preview_interval_min * 60
previews = model.get_previews()
static_previews = model.get_static_previews()
log_previews(iter, previews, static_previews)
if not is_reached_goal and (time.time() - last_save_time
) >= save_interval_min * 60:
last_save_time += save_interval_min * 60
model_save()
send_preview()
if i == 0:
if is_reached_goal:
model.pass_one_iter()
send_preview()
if debug:
time.sleep(0.005)
while not s2c.empty():
input = s2c.get()
op = input['op']
if op == 'save':
model_save()
elif op == 'backup':
model_backup()
elif op == 'preview':
if is_reached_goal:
model.pass_one_iter()
send_preview()
elif op == 'close':
model_save()
i = -1
break
if i == -1:
break
model.finalize()
except Exception as e:
print('Error: %s' % (str(e)))
traceback.print_exc()
break
c2s.put({'op': 'close'})
def ask_override(self):
return self.is_training and self.iter != 0 and io.input_in_time(
"两秒内按下Enter键可以重新设置模型参数...", 5 if io.is_colab() else 2)
