diff --git a/embedding_converter/README.md b/embedding_converter/README.md index 399f9db..2752f67 100644 --- a/embedding_converter/README.md +++ b/embedding_converter/README.md @@ -32,9 +32,9 @@ file_pattern = .datasets/megaface/**/*.jpg ``` [training.loader] -split_ratio = 0.95 batch_size = 256 num_workers = 8 +split_ratio = 0.95 ``` ``` diff --git a/embedding_converter/config.ini b/embedding_converter/config.ini index 5388075..e8360e2 100644 --- a/embedding_converter/config.ini +++ b/embedding_converter/config.ini @@ -2,9 +2,9 @@ file_pattern = [training.loader] -split_ratio = batch_size = num_workers = +split_ratio = [training.model] source_path = diff --git a/face_swapper/README.md b/face_swapper/README.md index 94943cd..518d481 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -28,7 +28,7 @@ This `config.ini` utilizes the MegaFace dataset to train the Face Swapper model. ``` [training.dataset] file_pattern = .datasets/vggface2/**/*.jpg -equal_probability = 0.2 +batch_ratio = 0.2 ``` ``` diff --git a/face_swapper/config.ini b/face_swapper/config.ini index c8a9f50..f578f27 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -1,6 +1,6 @@ [training.dataset] file_pattern = -equal_probability = +batch_ratio = [training.loader] batch_size = diff --git a/face_swapper/src/dataset.py b/face_swapper/src/dataset.py index e947de3..de7f344 100644 --- a/face_swapper/src/dataset.py +++ b/face_swapper/src/dataset.py @@ -10,15 +10,15 @@ from .types import Batch class DynamicDataset(Dataset[Tensor]): - def __init__(self, file_pattern : str, equal_probability : float) -> None: + def __init__(self, file_pattern : str, batch_ratio : float) -> None: self.file_paths = glob.glob(file_pattern) self.transforms = self.compose_transforms() - self.equal_probability = equal_probability + self.batch_ratio = batch_ratio def __getitem__(self, index : int) -> Batch: source_image_path = self.file_paths[index] - if random.random() < self.equal_probability: + if random.random() < self.batch_ratio: return self.prepare_equal_batch(source_image_path) return self.prepare_different_batch(source_image_path) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 20cfa77..9687048 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -143,10 +143,10 @@ def create_trainer() -> Trainer: def train() -> None: dataset_file_pattern = CONFIG.get('training.dataset', 'file_pattern') - dataset_equal_probability = CONFIG.getfloat('training.dataset', 'equal_probability') + dataset_batch_ratio = CONFIG.getfloat('training.dataset', 'batch_ratio') output_resume_path = CONFIG.get('training.output', 'resume_path') - dataset = DynamicDataset(dataset_file_pattern, dataset_equal_probability) + dataset = DynamicDataset(dataset_file_pattern, dataset_batch_ratio) training_loader, validation_loader = create_loaders(dataset) face_swapper_trainer = FaceSwapperTrainer() trainer = create_trainer()