From dc2b2dc982489638d1034e7e8293edd4a663d78e Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Fri, 11 Apr 2025 23:15:30 +0530 Subject: [PATCH 01/10] changes --- face_swapper/README.md | 1 + face_swapper/config.ini | 1 + face_swapper/src/models/loss.py | 23 ++++++++++++++++++++++- face_swapper/src/training.py | 10 ++++++++-- 4 files changed, 32 insertions(+), 3 deletions(-) diff --git a/face_swapper/README.md b/face_swapper/README.md index 98b8375..f55308d 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -77,6 +77,7 @@ num_filters = 16 ``` [training.losses] adversarial_weight = 1.0 +cycle_weight = 1.0 feature_weight = 10.0 reconstruction_weight = 10.0 identity_weight = 20.0 diff --git a/face_swapper/config.ini b/face_swapper/config.ini index 36ca46d..37f4627 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -37,6 +37,7 @@ num_filters = [training.losses] adversarial_weight = +cycle_weight = feature_weight = reconstruction_weight = identity_weight = diff --git a/face_swapper/src/models/loss.py b/face_swapper/src/models/loss.py index ef3a889..13841a2 100644 --- a/face_swapper/src/models/loss.py +++ b/face_swapper/src/models/loss.py @@ -49,7 +49,28 @@ class AdversarialLoss(nn.Module): return adversarial_loss, weighted_adversarial_loss -class FeatureLoss(nn.Module): +class CycleLoss(nn.Module): + def __init__(self, config_parser : ConfigParser) -> None: + super().__init__() + self.config_batch_size = config_parser.getint('training.loader', 'batch_size') + self.config_cycle_weight = config_parser.getfloat('training.losses', 'cycle_weight') + self.l1_loss = nn.L1Loss() + + def forward(self, target_tensor : Tensor, cycle_tensor : Tensor, target_features : Tuple[Feature, ...], cycle_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]: + temp_tensors = [] + + for target_feature, output_feature in zip(target_features, cycle_features): + temp_tensor = torch.mean(torch.pow(output_feature - target_feature, 2).reshape(self.config_batch_size, -1), dim = 1).mean() + temp_tensors.append(temp_tensor) + + cycle_feature_loss = torch.stack(temp_tensors).mean() + cycle_l1_loss = self.l1_loss(target_tensor, cycle_tensor) + cycle_loss = (cycle_feature_loss + cycle_l1_loss) * 0.5 + weighted_feature_loss = cycle_loss * self.config_cycle_weight + return cycle_loss, weighted_feature_loss + + +class FeautureLoss(nn.Module): def __init__(self, config_parser : ConfigParser) -> None: super().__init__() self.config_batch_size = config_parser.getint('training.loader', 'batch_size') diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 344fd4d..59c47f7 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -16,7 +16,7 @@ from .dataset import DynamicDataset from .helper import calc_embedding, overlay_mask from .models.discriminator import Discriminator from .models.generator import Generator -from .models.loss import AdversarialLoss, DiscriminatorLoss, FeatureLoss, GazeLoss, IdentityLoss, MaskLoss, MotionLoss, ReconstructionLoss +from .models.loss import AdversarialLoss, CycleLoss, DiscriminatorLoss, FeautureLoss, GazeLoss, IdentityLoss, MaskLoss, MotionLoss, ReconstructionLoss from .types import Batch, Embedding, Mask, OptimizerSet warnings.filterwarnings('ignore', category = UserWarning, module = 'torch') @@ -45,7 +45,8 @@ class FaceSwapperTrainer(LightningModule): self.discriminator = Discriminator(config_parser) self.discriminator_loss = DiscriminatorLoss() self.adversarial_loss = AdversarialLoss(config_parser) - self.feature_loss = FeatureLoss(config_parser) + self.cycle_loss = CycleLoss(config_parser) + self.feature_loss = FeautureLoss(config_parser) self.reconstruction_loss = ReconstructionLoss(config_parser, self.loss_embedder) self.identity_loss = IdentityLoss(config_parser, self.loss_embedder) self.motion_loss = MotionLoss(config_parser, self.motion_extractor) @@ -92,11 +93,15 @@ class FaceSwapperTrainer(LightningModule): generator_optimizer, discriminator_optimizer = self.optimizers() #type:ignore[attr-defined] source_embedding = calc_embedding(self.generator_embedder, source_tensor, (0, 0, 0, 0)) + target_embedding = calc_embedding(self.generator_embedder, target_tensor, (0, 0, 0, 0)) generator_target_features = self.generator.encode_features(target_tensor) generator_output_tensor, generator_output_mask = self.generator(source_embedding, target_tensor, generator_target_features) generator_output_features = self.generator.encode_features(generator_output_tensor) + cycle_output_tensor, cycle_output_mask = self.generator(target_embedding, generator_output_tensor, generator_output_features) + cycle_output_features = self.generator.encode_features(cycle_output_tensor) discriminator_output_tensors = self.discriminator(generator_output_tensor) adversarial_loss, weighted_adversarial_loss = self.adversarial_loss(discriminator_output_tensors) + cycle_loss, weighted_cycle_loss = self.cycle_loss(target_tensor, cycle_output_tensor, generator_target_features, cycle_output_features) feature_loss, weighted_feature_loss = self.feature_loss(generator_target_features, generator_output_features) reconstruction_loss, weighted_reconstruction_loss = self.reconstruction_loss(source_tensor, target_tensor, generator_output_tensor) identity_loss, weighted_identity_loss = self.identity_loss(generator_output_tensor, source_tensor) @@ -129,6 +134,7 @@ class FaceSwapperTrainer(LightningModule): self.log('generator_loss', generator_loss, prog_bar = True) self.log('discriminator_loss', discriminator_loss, prog_bar = True) self.log('adversarial_loss', adversarial_loss) + self.log('cycle_loss', cycle_loss) self.log('feature_loss', feature_loss) self.log('reconstruction_loss', reconstruction_loss) self.log('identity_loss', identity_loss) From b215db68c30141590dc5a5e2d3ea51fbb31c8635 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Fri, 11 Apr 2025 23:18:32 +0530 Subject: [PATCH 02/10] limit discriminator training every 10 steps --- face_swapper/src/training.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 59c47f7..ea03267 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -121,12 +121,13 @@ class FaceSwapperTrainer(LightningModule): generator_optimizer.zero_grad() self.untoggle_optimizer(generator_optimizer) - self.toggle_optimizer(discriminator_optimizer) - self.manual_backward(discriminator_loss) - if do_update: - discriminator_optimizer.step() - discriminator_optimizer.zero_grad() - self.untoggle_optimizer(discriminator_optimizer) + if self.global_step % 10 == 0: + self.toggle_optimizer(discriminator_optimizer) + self.manual_backward(discriminator_loss) + if do_update: + discriminator_optimizer.step() + discriminator_optimizer.zero_grad() + self.untoggle_optimizer(discriminator_optimizer) if self.global_step % self.config_preview_frequency == 0: self.generate_preview(source_tensor, target_tensor, generator_output_tensor, generator_output_mask) From b7a6f00e8bdfa0c433282ce29f5dd3889e3a1356 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 15 Apr 2025 13:36:21 +0530 Subject: [PATCH 03/10] remove discriminator frequency --- face_swapper/src/dataset.py | 4 ++-- face_swapper/src/training.py | 15 ++++++++------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/face_swapper/src/dataset.py b/face_swapper/src/dataset.py index 0a14eb4..db04932 100644 --- a/face_swapper/src/dataset.py +++ b/face_swapper/src/dataset.py @@ -88,13 +88,13 @@ class AugmentTransform: albumentations.OneOf( [ albumentations.MotionBlur(p = 0.1), - albumentations.ZoomBlur(p = 0.1) + albumentations.ZoomBlur(max_factor = (1.0, 1.1), p = 0.1) ], p = 0.2), albumentations.RandomBrightnessContrast(p = 0.7), albumentations.ColorJitter(p = 0.2), albumentations.RGBShift(p = 0.7), albumentations.Illumination(p = 0.2), - albumentations.Affine(translate_percent = (-0.03, 0.03), scale = (0.98, 1.02), rotate = (-2, 2), border_mode = 1, p = 0.7) + albumentations.Affine(translate_percent = (-0.03, 0.03), scale = (0.98, 1.02), rotate = (-2, 2), border_mode = 1, p = 0.3) ]) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index ea03267..1d994ae 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -116,18 +116,19 @@ class FaceSwapperTrainer(LightningModule): self.toggle_optimizer(generator_optimizer) self.manual_backward(generator_loss) + if do_update: generator_optimizer.step() generator_optimizer.zero_grad() self.untoggle_optimizer(generator_optimizer) - if self.global_step % 10 == 0: - self.toggle_optimizer(discriminator_optimizer) - self.manual_backward(discriminator_loss) - if do_update: - discriminator_optimizer.step() - discriminator_optimizer.zero_grad() - self.untoggle_optimizer(discriminator_optimizer) + self.toggle_optimizer(discriminator_optimizer) + self.manual_backward(discriminator_loss) + + if do_update: + discriminator_optimizer.step() + discriminator_optimizer.zero_grad() + self.untoggle_optimizer(discriminator_optimizer) if self.global_step % self.config_preview_frequency == 0: self.generate_preview(source_tensor, target_tensor, generator_output_tensor, generator_output_mask) From d9fe667cedf06000507e1222b131681154cea099 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Sun, 13 Apr 2025 10:17:52 +0200 Subject: [PATCH 04/10] Fix typo --- face_swapper/src/models/loss.py | 23 +---------------------- face_swapper/src/training.py | 12 ++---------- 2 files changed, 3 insertions(+), 32 deletions(-) diff --git a/face_swapper/src/models/loss.py b/face_swapper/src/models/loss.py index 13841a2..ef3a889 100644 --- a/face_swapper/src/models/loss.py +++ b/face_swapper/src/models/loss.py @@ -49,28 +49,7 @@ class AdversarialLoss(nn.Module): return adversarial_loss, weighted_adversarial_loss -class CycleLoss(nn.Module): - def __init__(self, config_parser : ConfigParser) -> None: - super().__init__() - self.config_batch_size = config_parser.getint('training.loader', 'batch_size') - self.config_cycle_weight = config_parser.getfloat('training.losses', 'cycle_weight') - self.l1_loss = nn.L1Loss() - - def forward(self, target_tensor : Tensor, cycle_tensor : Tensor, target_features : Tuple[Feature, ...], cycle_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]: - temp_tensors = [] - - for target_feature, output_feature in zip(target_features, cycle_features): - temp_tensor = torch.mean(torch.pow(output_feature - target_feature, 2).reshape(self.config_batch_size, -1), dim = 1).mean() - temp_tensors.append(temp_tensor) - - cycle_feature_loss = torch.stack(temp_tensors).mean() - cycle_l1_loss = self.l1_loss(target_tensor, cycle_tensor) - cycle_loss = (cycle_feature_loss + cycle_l1_loss) * 0.5 - weighted_feature_loss = cycle_loss * self.config_cycle_weight - return cycle_loss, weighted_feature_loss - - -class FeautureLoss(nn.Module): +class FeatureLoss(nn.Module): def __init__(self, config_parser : ConfigParser) -> None: super().__init__() self.config_batch_size = config_parser.getint('training.loader', 'batch_size') diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 1d994ae..344fd4d 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -16,7 +16,7 @@ from .dataset import DynamicDataset from .helper import calc_embedding, overlay_mask from .models.discriminator import Discriminator from .models.generator import Generator -from .models.loss import AdversarialLoss, CycleLoss, DiscriminatorLoss, FeautureLoss, GazeLoss, IdentityLoss, MaskLoss, MotionLoss, ReconstructionLoss +from .models.loss import AdversarialLoss, DiscriminatorLoss, FeatureLoss, GazeLoss, IdentityLoss, MaskLoss, MotionLoss, ReconstructionLoss from .types import Batch, Embedding, Mask, OptimizerSet warnings.filterwarnings('ignore', category = UserWarning, module = 'torch') @@ -45,8 +45,7 @@ class FaceSwapperTrainer(LightningModule): self.discriminator = Discriminator(config_parser) self.discriminator_loss = DiscriminatorLoss() self.adversarial_loss = AdversarialLoss(config_parser) - self.cycle_loss = CycleLoss(config_parser) - self.feature_loss = FeautureLoss(config_parser) + self.feature_loss = FeatureLoss(config_parser) self.reconstruction_loss = ReconstructionLoss(config_parser, self.loss_embedder) self.identity_loss = IdentityLoss(config_parser, self.loss_embedder) self.motion_loss = MotionLoss(config_parser, self.motion_extractor) @@ -93,15 +92,11 @@ class FaceSwapperTrainer(LightningModule): generator_optimizer, discriminator_optimizer = self.optimizers() #type:ignore[attr-defined] source_embedding = calc_embedding(self.generator_embedder, source_tensor, (0, 0, 0, 0)) - target_embedding = calc_embedding(self.generator_embedder, target_tensor, (0, 0, 0, 0)) generator_target_features = self.generator.encode_features(target_tensor) generator_output_tensor, generator_output_mask = self.generator(source_embedding, target_tensor, generator_target_features) generator_output_features = self.generator.encode_features(generator_output_tensor) - cycle_output_tensor, cycle_output_mask = self.generator(target_embedding, generator_output_tensor, generator_output_features) - cycle_output_features = self.generator.encode_features(cycle_output_tensor) discriminator_output_tensors = self.discriminator(generator_output_tensor) adversarial_loss, weighted_adversarial_loss = self.adversarial_loss(discriminator_output_tensors) - cycle_loss, weighted_cycle_loss = self.cycle_loss(target_tensor, cycle_output_tensor, generator_target_features, cycle_output_features) feature_loss, weighted_feature_loss = self.feature_loss(generator_target_features, generator_output_features) reconstruction_loss, weighted_reconstruction_loss = self.reconstruction_loss(source_tensor, target_tensor, generator_output_tensor) identity_loss, weighted_identity_loss = self.identity_loss(generator_output_tensor, source_tensor) @@ -116,7 +111,6 @@ class FaceSwapperTrainer(LightningModule): self.toggle_optimizer(generator_optimizer) self.manual_backward(generator_loss) - if do_update: generator_optimizer.step() generator_optimizer.zero_grad() @@ -124,7 +118,6 @@ class FaceSwapperTrainer(LightningModule): self.toggle_optimizer(discriminator_optimizer) self.manual_backward(discriminator_loss) - if do_update: discriminator_optimizer.step() discriminator_optimizer.zero_grad() @@ -136,7 +129,6 @@ class FaceSwapperTrainer(LightningModule): self.log('generator_loss', generator_loss, prog_bar = True) self.log('discriminator_loss', discriminator_loss, prog_bar = True) self.log('adversarial_loss', adversarial_loss) - self.log('cycle_loss', cycle_loss) self.log('feature_loss', feature_loss) self.log('reconstruction_loss', reconstruction_loss) self.log('identity_loss', identity_loss) From f4c4066e8c37996b30a18b3e14b6ac8915df2bb9 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Fri, 11 Apr 2025 23:15:30 +0530 Subject: [PATCH 05/10] changes --- face_swapper/src/models/loss.py | 21 +++++++++++++++++++++ face_swapper/src/training.py | 8 +++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/face_swapper/src/models/loss.py b/face_swapper/src/models/loss.py index ef3a889..2087a38 100644 --- a/face_swapper/src/models/loss.py +++ b/face_swapper/src/models/loss.py @@ -49,6 +49,27 @@ class AdversarialLoss(nn.Module): return adversarial_loss, weighted_adversarial_loss +class CycleLoss(nn.Module): + def __init__(self, config_parser : ConfigParser) -> None: + super().__init__() + self.config_batch_size = config_parser.getint('training.loader', 'batch_size') + self.config_cycle_weight = config_parser.getfloat('training.losses', 'cycle_weight') + self.l1_loss = nn.L1Loss() + + def forward(self, target_tensor : Tensor, cycle_tensor : Tensor, target_features : Tuple[Feature, ...], cycle_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]: + temp_tensors = [] + + for target_feature, output_feature in zip(target_features, cycle_features): + temp_tensor = torch.mean(torch.pow(output_feature - target_feature, 2).reshape(self.config_batch_size, -1), dim = 1).mean() + temp_tensors.append(temp_tensor) + + cycle_feature_loss = torch.stack(temp_tensors).mean() + cycle_l1_loss = self.l1_loss(target_tensor, cycle_tensor) + cycle_loss = (cycle_feature_loss + cycle_l1_loss) * 0.5 + weighted_feature_loss = cycle_loss * self.config_cycle_weight + return cycle_loss, weighted_feature_loss + + class FeatureLoss(nn.Module): def __init__(self, config_parser : ConfigParser) -> None: super().__init__() diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 344fd4d..183a46a 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -16,7 +16,7 @@ from .dataset import DynamicDataset from .helper import calc_embedding, overlay_mask from .models.discriminator import Discriminator from .models.generator import Generator -from .models.loss import AdversarialLoss, DiscriminatorLoss, FeatureLoss, GazeLoss, IdentityLoss, MaskLoss, MotionLoss, ReconstructionLoss +from .models.loss import AdversarialLoss, CycleLoss, DiscriminatorLoss, FeatureLoss, GazeLoss, IdentityLoss, MaskLoss, MotionLoss, ReconstructionLoss from .types import Batch, Embedding, Mask, OptimizerSet warnings.filterwarnings('ignore', category = UserWarning, module = 'torch') @@ -45,6 +45,7 @@ class FaceSwapperTrainer(LightningModule): self.discriminator = Discriminator(config_parser) self.discriminator_loss = DiscriminatorLoss() self.adversarial_loss = AdversarialLoss(config_parser) + self.cycle_loss = CycleLoss(config_parser) self.feature_loss = FeatureLoss(config_parser) self.reconstruction_loss = ReconstructionLoss(config_parser, self.loss_embedder) self.identity_loss = IdentityLoss(config_parser, self.loss_embedder) @@ -92,11 +93,15 @@ class FaceSwapperTrainer(LightningModule): generator_optimizer, discriminator_optimizer = self.optimizers() #type:ignore[attr-defined] source_embedding = calc_embedding(self.generator_embedder, source_tensor, (0, 0, 0, 0)) + target_embedding = calc_embedding(self.generator_embedder, target_tensor, (0, 0, 0, 0)) generator_target_features = self.generator.encode_features(target_tensor) generator_output_tensor, generator_output_mask = self.generator(source_embedding, target_tensor, generator_target_features) generator_output_features = self.generator.encode_features(generator_output_tensor) + cycle_output_tensor, cycle_output_mask = self.generator(target_embedding, generator_output_tensor, generator_output_features) + cycle_output_features = self.generator.encode_features(cycle_output_tensor) discriminator_output_tensors = self.discriminator(generator_output_tensor) adversarial_loss, weighted_adversarial_loss = self.adversarial_loss(discriminator_output_tensors) + cycle_loss, weighted_cycle_loss = self.cycle_loss(target_tensor, cycle_output_tensor, generator_target_features, cycle_output_features) feature_loss, weighted_feature_loss = self.feature_loss(generator_target_features, generator_output_features) reconstruction_loss, weighted_reconstruction_loss = self.reconstruction_loss(source_tensor, target_tensor, generator_output_tensor) identity_loss, weighted_identity_loss = self.identity_loss(generator_output_tensor, source_tensor) @@ -129,6 +134,7 @@ class FaceSwapperTrainer(LightningModule): self.log('generator_loss', generator_loss, prog_bar = True) self.log('discriminator_loss', discriminator_loss, prog_bar = True) self.log('adversarial_loss', adversarial_loss) + self.log('cycle_loss', cycle_loss) self.log('feature_loss', feature_loss) self.log('reconstruction_loss', reconstruction_loss) self.log('identity_loss', identity_loss) From 1477850a2340ae29fcfad965e4463426e14d6fa8 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Fri, 11 Apr 2025 23:18:32 +0530 Subject: [PATCH 06/10] limit discriminator training every 10 steps --- face_swapper/src/training.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 183a46a..554a457 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -121,12 +121,13 @@ class FaceSwapperTrainer(LightningModule): generator_optimizer.zero_grad() self.untoggle_optimizer(generator_optimizer) - self.toggle_optimizer(discriminator_optimizer) - self.manual_backward(discriminator_loss) - if do_update: - discriminator_optimizer.step() - discriminator_optimizer.zero_grad() - self.untoggle_optimizer(discriminator_optimizer) + if self.global_step % 10 == 0: + self.toggle_optimizer(discriminator_optimizer) + self.manual_backward(discriminator_loss) + if do_update: + discriminator_optimizer.step() + discriminator_optimizer.zero_grad() + self.untoggle_optimizer(discriminator_optimizer) if self.global_step % self.config_preview_frequency == 0: self.generate_preview(source_tensor, target_tensor, generator_output_tensor, generator_output_mask) From 39ce14b5903ce117874d23d0739989b5c4c60e20 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 15 Apr 2025 13:36:21 +0530 Subject: [PATCH 07/10] remove discriminator frequency --- face_swapper/src/training.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 554a457..8d1d749 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -116,18 +116,19 @@ class FaceSwapperTrainer(LightningModule): self.toggle_optimizer(generator_optimizer) self.manual_backward(generator_loss) + if do_update: generator_optimizer.step() generator_optimizer.zero_grad() self.untoggle_optimizer(generator_optimizer) - if self.global_step % 10 == 0: - self.toggle_optimizer(discriminator_optimizer) - self.manual_backward(discriminator_loss) - if do_update: - discriminator_optimizer.step() - discriminator_optimizer.zero_grad() - self.untoggle_optimizer(discriminator_optimizer) + self.toggle_optimizer(discriminator_optimizer) + self.manual_backward(discriminator_loss) + + if do_update: + discriminator_optimizer.step() + discriminator_optimizer.zero_grad() + self.untoggle_optimizer(discriminator_optimizer) if self.global_step % self.config_preview_frequency == 0: self.generate_preview(source_tensor, target_tensor, generator_output_tensor, generator_output_mask) From 4a319ec9bd0552451a326f4ff617d66e46471fa3 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 15 Apr 2025 14:05:27 +0530 Subject: [PATCH 08/10] changes --- face_swapper/src/models/loss.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/face_swapper/src/models/loss.py b/face_swapper/src/models/loss.py index 2087a38..ccbd96a 100644 --- a/face_swapper/src/models/loss.py +++ b/face_swapper/src/models/loss.py @@ -54,7 +54,7 @@ class CycleLoss(nn.Module): super().__init__() self.config_batch_size = config_parser.getint('training.loader', 'batch_size') self.config_cycle_weight = config_parser.getfloat('training.losses', 'cycle_weight') - self.l1_loss = nn.L1Loss() + self.mae_loss = nn.L1Loss() def forward(self, target_tensor : Tensor, cycle_tensor : Tensor, target_features : Tuple[Feature, ...], cycle_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]: temp_tensors = [] @@ -63,9 +63,9 @@ class CycleLoss(nn.Module): temp_tensor = torch.mean(torch.pow(output_feature - target_feature, 2).reshape(self.config_batch_size, -1), dim = 1).mean() temp_tensors.append(temp_tensor) - cycle_feature_loss = torch.stack(temp_tensors).mean() - cycle_l1_loss = self.l1_loss(target_tensor, cycle_tensor) - cycle_loss = (cycle_feature_loss + cycle_l1_loss) * 0.5 + feature_loss = torch.stack(temp_tensors).mean() + mae_loss = self.mae_loss(target_tensor, cycle_tensor) + cycle_loss = (feature_loss + mae_loss) * 0.5 weighted_feature_loss = cycle_loss * self.config_cycle_weight return cycle_loss, weighted_feature_loss From 128726701b4bd5c5f0e5f7b6af0983dc6aef7e9b Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 15 Apr 2025 14:10:58 +0530 Subject: [PATCH 09/10] changes --- face_swapper/src/models/loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/face_swapper/src/models/loss.py b/face_swapper/src/models/loss.py index ccbd96a..1cca867 100644 --- a/face_swapper/src/models/loss.py +++ b/face_swapper/src/models/loss.py @@ -64,8 +64,8 @@ class CycleLoss(nn.Module): temp_tensors.append(temp_tensor) feature_loss = torch.stack(temp_tensors).mean() - mae_loss = self.mae_loss(target_tensor, cycle_tensor) - cycle_loss = (feature_loss + mae_loss) * 0.5 + reconstruction_loss = self.mae_loss(target_tensor, cycle_tensor) + cycle_loss = (feature_loss + reconstruction_loss) * 0.5 weighted_feature_loss = cycle_loss * self.config_cycle_weight return cycle_loss, weighted_feature_loss From bcf5b4e5a86f6f471900543f910c7192186daeef Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 15 Apr 2025 14:14:09 +0530 Subject: [PATCH 10/10] changes --- face_swapper/src/models/loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/face_swapper/src/models/loss.py b/face_swapper/src/models/loss.py index 1cca867..82642bb 100644 --- a/face_swapper/src/models/loss.py +++ b/face_swapper/src/models/loss.py @@ -54,7 +54,7 @@ class CycleLoss(nn.Module): super().__init__() self.config_batch_size = config_parser.getint('training.loader', 'batch_size') self.config_cycle_weight = config_parser.getfloat('training.losses', 'cycle_weight') - self.mae_loss = nn.L1Loss() + self.l1_loss = nn.L1Loss() def forward(self, target_tensor : Tensor, cycle_tensor : Tensor, target_features : Tuple[Feature, ...], cycle_features : Tuple[Feature, ...]) -> Tuple[Loss, Loss]: temp_tensors = [] @@ -64,7 +64,7 @@ class CycleLoss(nn.Module): temp_tensors.append(temp_tensor) feature_loss = torch.stack(temp_tensors).mean() - reconstruction_loss = self.mae_loss(target_tensor, cycle_tensor) + reconstruction_loss = self.l1_loss(target_tensor, cycle_tensor) cycle_loss = (feature_loss + reconstruction_loss) * 0.5 weighted_feature_loss = cycle_loss * self.config_cycle_weight return cycle_loss, weighted_feature_loss