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/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/models/loss.py b/face_swapper/src/models/loss.py index ef3a889..82642bb 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) + + feature_loss = torch.stack(temp_tensors).mean() + 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 + + 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..8d1d749 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) @@ -111,6 +116,7 @@ class FaceSwapperTrainer(LightningModule): self.toggle_optimizer(generator_optimizer) self.manual_backward(generator_loss) + if do_update: generator_optimizer.step() generator_optimizer.zero_grad() @@ -118,6 +124,7 @@ class FaceSwapperTrainer(LightningModule): self.toggle_optimizer(discriminator_optimizer) self.manual_backward(discriminator_loss) + if do_update: discriminator_optimizer.step() discriminator_optimizer.zero_grad() @@ -129,6 +136,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)