diff --git a/hyperswap/README.md b/hyperswap/README.md index 1f6b7b3..4022916 100644 --- a/hyperswap/README.md +++ b/hyperswap/README.md @@ -47,7 +47,6 @@ split_ratio = 0.9995 [training.model] generator_embedder_path = .models/blendface.pt loss_embedder_path = .models/arcface.pt -gazer_path = .models/gazer.pt face_masker_path = .models/face_masker.pt ``` @@ -77,11 +76,9 @@ 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 -gaze_weight = 0.05 mask_weight = 5.0 ``` diff --git a/hyperswap/config.ini b/hyperswap/config.ini index f4bea66..7eee9e6 100644 --- a/hyperswap/config.ini +++ b/hyperswap/config.ini @@ -15,7 +15,6 @@ split_ratio = [training.model] generator_embedder_path = loss_embedder_path = -gazer_path = face_masker_path = [training.model.generator] @@ -37,11 +36,9 @@ num_filters = [training.losses] adversarial_weight = -cycle_weight = feature_weight = reconstruction_weight = identity_weight = -gaze_weight = mask_weight = [training.trainer] diff --git a/hyperswap/src/models/loss.py b/hyperswap/src/models/loss.py index 06552de..c159de9 100644 --- a/hyperswap/src/models/loss.py +++ b/hyperswap/src/models/loss.py @@ -4,10 +4,9 @@ from typing import List, Tuple import torch from pytorch_msssim import ssim from torch import Tensor, nn -from torchvision import transforms from ..helper import calculate_face_embedding, dilate_mask -from ..types import EmbedderModule, FaceMaskerModule, Feature, GazerModule, Loss, Mask +from ..types import EmbedderModule, FaceMaskerModule, Feature, Loss, Mask class DiscriminatorLoss(nn.Module): @@ -49,27 +48,6 @@ 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__() @@ -127,38 +105,6 @@ class IdentityLoss(nn.Module): return identity_loss, weighted_identity_loss -class GazeLoss(nn.Module): - def __init__(self, config_parser : ConfigParser, gazer : GazerModule) -> None: - super().__init__() - self.config_gaze_weight = config_parser.getfloat('training.losses', 'gaze_weight') - self.config_output_size = config_parser.getint('training.model.generator', 'output_size') - self.gazer = gazer - self.l1_loss = nn.L1Loss() - - def forward(self, target_tensor : Tensor, output_tensor : Tensor) -> Tuple[Loss, Loss]: - output_pitch, output_yaw = self.detect_gaze(output_tensor) - target_pitch, target_yaw = self.detect_gaze(target_tensor) - - pitch_loss = self.l1_loss(output_pitch, target_pitch) - yaw_loss = self.l1_loss(output_yaw, target_yaw) - - gaze_loss = (pitch_loss + yaw_loss) * 0.5 - weighted_gaze_loss = gaze_loss * self.config_gaze_weight - return gaze_loss, weighted_gaze_loss - - def detect_gaze(self, input_tensor : Tensor) -> Tuple[Tensor, Tensor]: - crop_sizes = (torch.tensor([ 0.235, 0.875, 0.0625, 0.8 ]) * self.config_output_size).int() - crop_tensor = input_tensor[:, :, crop_sizes[0]:crop_sizes[1], crop_sizes[2]:crop_sizes[3]] - crop_tensor = (crop_tensor + 1) * 0.5 - crop_tensor = transforms.Normalize(mean = [ 0.485, 0.456, 0.406 ], std = [ 0.229, 0.224, 0.225 ])(crop_tensor) - crop_tensor = nn.functional.interpolate(crop_tensor, size = 448, mode = 'bicubic') - - with torch.no_grad(): - pitch, yaw = self.gazer(crop_tensor) - - return pitch, yaw - - class MaskLoss(nn.Module): def __init__(self, config_parser : ConfigParser, face_masker : FaceMaskerModule) -> None: super().__init__() diff --git a/hyperswap/src/training.py b/hyperswap/src/training.py index 7b89366..7d52478 100644 --- a/hyperswap/src/training.py +++ b/hyperswap/src/training.py @@ -19,7 +19,7 @@ from .dataset import DynamicDataset from .helper import apply_noise, calculate_face_embedding, erode_mask, overlay_mask from .models.discriminator import Discriminator from .models.generator import Generator -from .models.loss import AdversarialLoss, CycleLoss, DiscriminatorLoss, FeatureLoss, GazeLoss, IdentityLoss, MaskLoss, ReconstructionLoss +from .models.loss import AdversarialLoss, DiscriminatorLoss, FeatureLoss, IdentityLoss, MaskLoss, ReconstructionLoss from .types import Batch, Embedding, Mask, OptimizerSet, TrainerPrecision, TrainerStrategy warnings.filterwarnings('ignore', category = UserWarning, module = 'torch') @@ -33,7 +33,6 @@ class HyperSwapTrainer(LightningModule): super().__init__() self.config_generator_embedder_path = config_parser.get('training.model', 'generator_embedder_path') self.config_loss_embedder_path = config_parser.get('training.model', 'loss_embedder_path') - self.config_gazer_path = config_parser.get('training.model', 'gazer_path') self.config_face_masker_path = config_parser.get('training.model', 'face_masker_path') self.config_accumulate_size = config_parser.getfloat('training.trainer', 'accumulate_size') self.config_discriminator_ratio = config_parser.getfloat('training.trainer', 'discriminator_ratio') @@ -51,17 +50,14 @@ class HyperSwapTrainer(LightningModule): self.config_discriminator_scheduler_patience = config_parser.getint('training.optimizer.discriminator', 'scheduler_patience') self.generator_embedder = torch.jit.load(self.config_generator_embedder_path, map_location = 'cpu').eval() self.loss_embedder = torch.jit.load(self.config_loss_embedder_path, map_location = 'cpu').eval() - self.gazer = torch.jit.load(self.config_gazer_path, map_location = 'cpu').eval() self.face_masker = torch.jit.load(self.config_face_masker_path, map_location ='cpu').eval() self.generator = Generator(config_parser) 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) - self.gaze_loss = GazeLoss(config_parser, self.gazer) self.mask_loss = MaskLoss(config_parser, self.face_masker) self.automatic_optimization = False @@ -105,7 +101,6 @@ class HyperSwapTrainer(LightningModule): generator_optimizer, discriminator_optimizer = self.optimizers() #type:ignore[attr-defined] generator_scheduler, discriminator_scheduler = self.lr_schedulers() #type:ignore[attr-defined] source_embedding = calculate_face_embedding(self.generator_embedder, source_tensor, (0, 0, 0, 0)) - target_embedding = calculate_face_embedding(self.generator_embedder, target_tensor, (0, 0, 0, 0)) if self.config_noise_factor > 0: source_embedding = apply_noise(source_embedding, self.config_noise_factor) @@ -114,17 +109,13 @@ class HyperSwapTrainer(LightningModule): 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) - gaze_loss, weighted_gaze_loss = self.gaze_loss(target_tensor, generator_output_tensor) mask_loss, weighted_mask_loss = self.mask_loss(target_tensor, generator_output_mask) - generator_loss = weighted_adversarial_loss + weighted_cycle_loss + weighted_feature_loss + weighted_reconstruction_loss + weighted_identity_loss + weighted_gaze_loss + weighted_mask_loss + generator_loss = weighted_adversarial_loss + weighted_feature_loss + weighted_reconstruction_loss + weighted_identity_loss + weighted_mask_loss if torch.randn(1).item() < self.config_discriminator_ratio: discriminator_real_tensors = self.discriminator(source_tensor) @@ -167,11 +158,9 @@ class HyperSwapTrainer(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) - self.log('gaze_loss', gaze_loss) self.log('mask_loss', mask_loss) if do_update: diff --git a/hyperswap/src/types.py b/hyperswap/src/types.py index ffc597c..1f941ec 100644 --- a/hyperswap/src/types.py +++ b/hyperswap/src/types.py @@ -19,7 +19,6 @@ Padding : TypeAlias = Tuple[int, int, int, int] GeneratorModule : TypeAlias = Module EmbedderModule : TypeAlias = Module -GazerModule : TypeAlias = Module FaceMaskerModule : TypeAlias = Module OptimizerSet : TypeAlias = Any