github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-08-21 22:13:20 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

support run cpu trainer without preview

Browse source
This commit is contained in:
plucky 2018-12-28 16:03:48 +08:00
commit cacacfce9a
3 changed files with 8 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

4
cpu.sh
View file

@ -47,7 +47,9 @@ do
;;
"train" )
echo -n "Model? [ H64 (2GB+) | H128 (3GB+) | DF (5GB+) | LIAEF128 (5GB+) | LIAEF128YAW (5GB+) | MIAEF128 (5GB+) | AVATAR (4GB+) ] "; read model
$PYTHON $INTERNAL_DIR/main.py train --training-data-src-dir $WORKSPACE/data_src/aligned --training-data-dst-dir $WORKSPACE/data_dst/aligned --model-dir $WORKSPACE/model --model $model --cpu-only
echo -n "Show Preview? [Y/n] "; read preview
if [ "$preview" == "Y" ] || [ "$preview" == "y" ]; then preview="--preview"; else preview=""; fi
$PYTHON $INTERNAL_DIR/main.py train --training-data-src-dir $WORKSPACE/data_src/aligned --training-data-dst-dir $WORKSPACE/data_dst/aligned --model-dir $WORKSPACE/model --model $model --cpu-only $preview
;;
"convert" )
echo -n "Model? [ H64 (2GB+) | H128 (3GB+) | DF (5GB+) | LIAEF128 (5GB+) | LIAEF128YAW (5GB+) | MIAEF128 (5GB+) | AVATAR(4GB+) ] "; read model

5
main.py
View file

@ -66,13 +66,11 @@ if __name__ == "__main__":
sort_parser.set_defaults (func=process_sort)
def process_train(arguments):
if 'DFL_TARGET_EPOCH' in os.environ.keys():
arguments.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'] )
from mainscripts import Trainer
Trainer.main (
training_data_src_dir=arguments.training_data_src_dir,
@ -80,6 +78,7 @@ if __name__ == "__main__":
model_path=arguments.model_dir,
model_name=arguments.model_name,
debug = arguments.debug,
preview = arguments.preview,
#**options
batch_size = arguments.batch_size,
write_preview_history = arguments.write_preview_history,
@ -99,7 +98,6 @@ 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('--write-preview-history', action="store_true", dest="write_preview_history", default=False, help="Enable write preview history.")
train_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug training.")
train_parser.add_argument('--preview', action="store_true", dest="preview", default=True, help="Show preview.")
train_parser.add_argument('--batch-size', type=int, dest="batch_size", default=0, help="Model batch size. Default - auto. Environment variable: ODFS_BATCH_SIZE.")
train_parser.add_argument('--target-epoch', type=int, dest="target_epoch", default=0, help="Train until target epoch. Default - unlimited. Environment variable: ODFS_TARGET_EPOCH.")
train_parser.add_argument('--save-interval-min', type=int, dest="save_interval_min", default=10, help="Save interval in minutes. Default 10.")
@ -108,6 +106,7 @@ if __name__ == "__main__":
train_parser.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="MultiGPU option. It will select only same best(worst) GPU models.")
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('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
train_parser.add_argument('--preview', action="store_true",dest="preview", default=False, help="Show preview.")
train_parser.set_defaults (func=process_train)

5
mainscripts/Trainer.py
View file

@ -102,6 +102,7 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
model.pass_one_epoch()
send_preview()
elif op == 'close':
print ('Closing and Saving.')
model_save()
i = -1
break
@ -278,12 +279,12 @@ def previewThread (input_queue, output_queue):
cv2.destroyAllWindows()
def main (training_data_src_dir, training_data_dst_dir, model_path, model_name,preview, **in_options):
print ("Running trainer.\r\n")
print ("Running trainer(preview=%s).\r\n" % (preview))
output_queue = queue.Queue()
input_queue = queue.Queue()
import threading
thread = threading.Thread(target=trainerThread, args=(output_queue, input_queue, training_data_src_dir, training_data_dst_dir, model_path, model_name), kwargs=in_options )
thread.start()
if preview:
previewThread (input_queue, output_queue)