CalvinBackup/neuralchen-SimSwap
1
0
Fork 0
Code Issues 288 Pull Requests 16 Actions Packages Projects Releases 2 Wiki Activity

Update train.py

Browse Source
This commit is contained in:
chenxuanhong
2022-04-21 22:56:03 +08:00
parent 7c44bc4b9a
commit bc5ac1ef22
1 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+9 -7
View File
@@ -5,7 +5,7 @@
# Created Date: Monday December 27th 2021 # Created Date: Monday December 27th 2021
# Author: Chen Xuanhong # Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com # Email: chenxuanhongzju@outlook.com
# Last Modified: Thursday, 21st April 2022 10:36:48 pm # Last Modified: Thursday, 21st April 2022 10:53:51 pm
# Modified By: Chen Xuanhong # Modified By: Chen Xuanhong
# Copyright (c) 2021 Shanghai Jiao Tong University # Copyright (c) 2021 Shanghai Jiao Tong University
############################################################# #############################################################
@@ -48,6 +48,7 @@ class TrainOptions:
# for displays # for displays
self.parser.add_argument('--tag', type=str, default='simswap') self.parser.add_argument('--tag', type=str, default='simswap')
self.parser.add_argument('--use_tensorboard', type=str2bool, default='True')
# for training # for training
self.parser.add_argument('--dataset', type=str, default="G:/VGGFace2-HQ/VGGface2_None_norm_512_true_bygfpgan", help='path to the face swapping dataset') self.parser.add_argument('--dataset', type=str, default="G:/VGGFace2-HQ/VGGface2_None_norm_512_true_bygfpgan", help='path to the face swapping dataset')
@@ -142,9 +143,10 @@ if __name__ == '__main__':
model.initialize(opt) model.initialize(opt)
##################################################### #####################################################
if opt.use_tensorboard:
tensorboard_writer = tensorboard.SummaryWriter(log_path) tensorboard_writer = tensorboard.SummaryWriter(log_path)
logger = tensorboard_writer logger = tensorboard_writer
log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt') log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
with open(log_name, "a") as log_file: with open(log_name, "a") as log_file:
@@ -247,9 +249,9 @@ if __name__ == '__main__':
"D_real":loss_Dreal.item(), "D_real":loss_Dreal.item(),
"D_loss":loss_D.item() "D_loss":loss_D.item()
} }
if opt.use_tensorboard:
for tag, value in errors.items(): for tag, value in errors.items():
logger.add_scalar(tag, value, step) logger.add_scalar(tag, value, step)
message = '( step: %d, ) ' % (step) message = '( step: %d, ) ' % (step)
for k, v in errors.items(): for k, v in errors.items():
message += '%s: %.3f ' % (k, v) message += '%s: %.3f ' % (k, v)