From a5eb7d6aa1794e680913814f5866137273d5eb2f Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Thu, 27 Feb 2025 19:05:41 +0530 Subject: [PATCH] change Adam to AdamW --- face_swapper/src/training.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 01cb2c8..704be44 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -46,8 +46,8 @@ class FaceSwapperTrainer(lightning.LightningModule): def configure_optimizers(self) -> Tuple[OptimizerConfig, OptimizerConfig]: learning_rate = CONFIG.getfloat('training.trainer', 'learning_rate') - generator_optimizer = torch.optim.Adam(self.generator.parameters(), lr = learning_rate, betas = (0.0, 0.999), weight_decay = 1e-4) - discriminator_optimizer = torch.optim.Adam(self.discriminator.parameters(), lr = learning_rate, betas = (0.0, 0.999), weight_decay = 1e-4) + generator_optimizer = torch.optim.AdamW(self.generator.parameters(), lr = learning_rate, betas = (0.0, 0.999), weight_decay = 1e-4) + discriminator_optimizer = torch.optim.AdamW(self.discriminator.parameters(), lr = learning_rate, betas = (0.0, 0.999), weight_decay = 1e-4) generator_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(generator_optimizer, T_0 = 300, T_mult = 2) discriminator_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(discriminator_optimizer, T_0 = 300, T_mult = 2)