feat: optimization eps value control

jolibrain · Oct 11, 2023 · 0556987 · 0556987
1 parent 124ac0f
commit 0556987
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 4 deletions.
diff --git a/models/base_model.py b/models/base_model.py
@@ -221,6 +221,7 @@ def init_semantic_cls(self, opt):
                 lr=opt.train_sem_lr_f_s,
                 betas=(opt.train_beta1, opt.train_beta2),
                 weight_decay=opt.train_optim_weight_decay,
+                eps=opt.train_optim_eps,
             )
 
             if opt.train_cls_regression:
@@ -323,6 +324,7 @@ def init_semantic_mask(self, opt):
                         lr=opt.train_sem_lr_f_s,
                         betas=(opt.train_beta1, opt.train_beta2),
                         weight_decay=opt.train_optim_weight_decay,
+                        eps=opt.train_optim_eps,
                     )
                 else:
                     self.optimizer_f_s = opt.optim(
@@ -331,6 +333,7 @@ def init_semantic_mask(self, opt):
                         lr=opt.train_sem_lr_f_s,
                         betas=(opt.train_beta1, opt.train_beta2),
                         weight_decay=opt.train_optim_weight_decay,
+                        eps=opt.train_optim_eps,
                     )
 
                 self.optimizers.append(self.optimizer_f_s)

diff --git a/models/cut_model.py b/models/cut_model.py
@@ -254,6 +254,7 @@ def __init__(self, opt, rank):
                 lr=opt.train_G_lr,
                 betas=(opt.train_beta1, opt.train_beta2),
                 weight_decay=opt.train_optim_weight_decay,
+                eps=opt.train_optim_eps,
             )
             if self.opt.model_multimodal:
                 self.criterionZ = torch.nn.L1Loss()
@@ -263,6 +264,7 @@ def __init__(self, opt, rank):
                     lr=opt.train_G_lr,
                     betas=(opt.train_beta1, opt.train_beta2),
                     weight_decay=opt.train_optim_weight_decay,
+                    eps=opt.train_optim_eps,
                 )
 
             if len(self.discriminators_names) > 0:
@@ -283,6 +285,7 @@ def __init__(self, opt, rank):
                 lr=opt.train_D_lr,
                 betas=(opt.train_beta1, opt.train_beta2),
                 weight_decay=opt.train_optim_weight_decay,
+                eps=opt.train_optim_eps,
             )
 
             self.optimizers.append(self.optimizer_G)
@@ -426,6 +429,7 @@ def data_dependent_initialize(self, data):
                     lr=self.opt.train_G_lr,
                     betas=(self.opt.train_beta1, self.opt.train_beta2),
                     weight_decay=self.opt.train_optim_weight_decay,
+                    eps=self.opt.train_optim_eps,
                 )
                 self.optimizers.append(self.optimizer_F)
 

diff --git a/models/cycle_gan_model.py b/models/cycle_gan_model.py
@@ -136,6 +136,7 @@ def __init__(self, opt, rank):
                 lr=opt.train_G_lr,
                 betas=(opt.train_beta1, opt.train_beta2),
                 weight_decay=opt.train_optim_weight_decay,
+                eps=opt.train_optim_eps,
             )
 
             D_parameters = itertools.chain(
@@ -151,6 +152,7 @@ def __init__(self, opt, rank):
                 lr=opt.train_D_lr,
                 betas=(opt.train_beta1, opt.train_beta2),
                 weight_decay=opt.train_optim_weight_decay,
+                eps=opt.train_optim_eps,
             )
 
             self.optimizers.append(self.optimizer_G)

diff --git a/models/palette_model.py b/models/palette_model.py
@@ -344,6 +344,7 @@ def __init__(self, opt, rank):
                 lr=opt.train_G_lr,
                 betas=(opt.train_beta1, opt.train_beta2),
                 weight_decay=opt.train_optim_weight_decay,
+                eps=opt.train_optim_eps,
             )
             self.optimizers.append(self.optimizer_G)
 

diff --git a/options/train_options.py b/options/train_options.py
@@ -179,6 +179,12 @@ def initialize(self, parser):
             default=0.0,
             help="weight decay for optimizer",
         )
+        parser.add_argument(
+            "--train_optim_eps",
+            type=float,
+            default=1e-8,
+            help="epsilon for optimizer",
+        )
         parser.add_argument(
             "--train_load_iter",
             type=int,

diff --git a/train.py b/train.py
@@ -51,16 +51,16 @@ def setup(rank, world_size, port):
     dist.init_process_group("nccl", rank=rank, world_size=world_size)
 
 
-def optim(opt, params, lr, betas, weight_decay):
+def optim(opt, params, lr, betas, weight_decay, eps):
     print("Using ", opt.train_optim, " as optimizer")
     if opt.train_optim == "adam":
-        return torch.optim.Adam(params, lr, betas, weight_decay=weight_decay)
+        return torch.optim.Adam(params, lr, betas, weight_decay=weight_decay, eps=eps)
     elif opt.train_optim == "radam":
-        return torch.optim.RAdam(params, lr, betas, weight_decay=weight_decay)
+        return torch.optim.RAdam(params, lr, betas, weight_decay=weight_decay, eps=eps)
     elif opt.train_optim == "adamw":
         if weight_decay == 0.0:
             weight_decay = 0.01  # default value
-        return torch.optim.AdamW(params, lr, betas, weight_decay=weight_decay)
+        return torch.optim.AdamW(params, lr, betas, weight_decay=weight_decay, eps=eps)
     elif opt.train_optim == "lion":
         return Lion(params, lr, betas, weight_decay)