diff --git a/embedding_converter/README.md b/embedding_converter/README.md index 502670f..3f940c1 100644 --- a/embedding_converter/README.md +++ b/embedding_converter/README.md @@ -47,7 +47,10 @@ target_path = .models/arcface_simswap.pt [training.trainer] learning_rate = 0.001 max_epochs = 4096 +strategy = auto precision = 16-mixed +logger_path = .logs +logger_name = arcface_converter_simswap ``` ``` diff --git a/embedding_converter/config.ini b/embedding_converter/config.ini index cb24d3f..a040687 100644 --- a/embedding_converter/config.ini +++ b/embedding_converter/config.ini @@ -13,7 +13,10 @@ target_path = [training.trainer] learning_rate = max_epochs = +strategy = precision = +logger_path = +logger_name = [training.output] directory_path = diff --git a/embedding_converter/src/training.py b/embedding_converter/src/training.py index 50d42d2..8333fc8 100644 --- a/embedding_converter/src/training.py +++ b/embedding_converter/src/training.py @@ -89,15 +89,19 @@ def split_dataset(dataset : Dataset[Tensor]) -> Tuple[Dataset[Tensor], Dataset[T def create_trainer() -> Trainer: config_max_epochs = CONFIG_PARSER.getint('training.trainer', 'max_epochs') + config_strategy = CONFIG_PARSER.get('training.trainer', 'strategy') config_precision = CONFIG_PARSER.get('training.trainer', 'precision') + config_logger_path = CONFIG_PARSER.get('training.trainer', 'logger_path') + config_logger_name = CONFIG_PARSER.get('training.trainer', 'logger_name') config_directory_path = CONFIG_PARSER.get('training.output', 'directory_path') config_file_pattern = CONFIG_PARSER.get('training.output', 'file_pattern') - logger = TensorBoardLogger('.logs', name = 'embedding_converter') + logger = TensorBoardLogger(config_logger_path, config_logger_name) return Trainer( logger = logger, log_every_n_steps = 10, max_epochs = config_max_epochs, + strategy = config_strategy, precision = config_precision, callbacks = [ diff --git a/face_swapper/README.md b/face_swapper/README.md index 5776b31..fa85087 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -89,9 +89,11 @@ expression_weight = 0.0 accumulate_size = 4 learning_rate = 0.0004 max_epochs = 50 -precision = 16-mixed -preview_frequency = 250 strategy = auto +precision = 16-mixed +logger_path = .logs +logger_name = face_swapper +preview_frequency = 250 ``` ``` diff --git a/face_swapper/config.ini b/face_swapper/config.ini index 1ff4100..79ac834 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -47,9 +47,11 @@ expression_weight = accumulate_size = learning_rate = max_epochs = -precision = -preview_frequency = strategy = +precision = +logger_path = +logger_name = +preview_frequency = [training.output] directory_path = diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 68a0aad..f5a62b0 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -208,18 +208,20 @@ def split_dataset(dataset : Dataset[Tensor]) -> Tuple[Dataset[Tensor], Dataset[T def create_trainer() -> Trainer: config_max_epochs = CONFIG_PARSER.getint('training.trainer', 'max_epochs') - config_precision = CONFIG_PARSER.get('training.trainer', 'precision') config_strategy = CONFIG_PARSER.get('training.trainer', 'strategy') + config_precision = CONFIG_PARSER.get('training.trainer', 'precision') + config_logger_path = CONFIG_PARSER.get('training.trainer', 'logger_path') + config_logger_name = CONFIG_PARSER.get('training.trainer', 'logger_name') config_directory_path = CONFIG_PARSER.get('training.output', 'directory_path') config_file_pattern = CONFIG_PARSER.get('training.output', 'file_pattern') - logger = TensorBoardLogger('.logs', name = 'face_swapper') + logger = TensorBoardLogger(config_logger_path, config_logger_name) return Trainer( logger = logger, log_every_n_steps = 10, max_epochs = config_max_epochs, - precision = config_precision, strategy = config_strategy, + precision = config_precision, callbacks = [ ModelCheckpoint(