x-evolution 0.1.11__tar.gz → 0.1.15__tar.gz

{x_evolution-0.1.11 → x_evolution-0.1.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.11
+Version: 0.1.15
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution

{x_evolution-0.1.11 → x_evolution-0.1.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-evolution"
-version = "0.1.11"
+version = "0.1.15"
 description = "x-evolution"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_evolution-0.1.11 → x_evolution-0.1.15}/train_lunar.py

@@ -80,6 +80,8 @@ from x_mlps_pytorch.residual_normed_mlp import ResidualNormedMLP
 
 actor = ResidualNormedMLP(dim_in = 8, dim = 24, depth = 2, residual_every = 1, dim_out = 4)
 
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
 evo_strat = EvoStrategy(
     actor,
     environment = LunarEnvironment(repeats = 2),
@@ -91,7 +93,10 @@ evo_strat = EvoStrategy(
     learned_noise_scale = True,
     use_sigma_optimizer = True,
     learning_rate = 1e-3,
-    noise_scale_learning_rate = 1e-5
+    noise_scale_learning_rate = 1e-4,
+    use_scheduler = True,
+    scheduler_klass = CosineAnnealingLR,
+    scheduler_kwargs = dict(T_max = 50_000)
 )
 
 evo_strat()

{x_evolution-0.1.11 → x_evolution-0.1.15}/train_mnist.py

@@ -52,12 +52,12 @@ evo_strat = EvoStrategy(
     environment = loss_mnist,
     noise_population_size = 100,
     noise_scale = 1e-2,
-    noise_scale_clamp_range = (5e-3, 2e-2),
-    noise_low_rank = 2,
+    noise_scale_clamp_range = (8e-3, 2e-2),
+    noise_low_rank = 1,
     num_generations = 10_000,
     learning_rate = 1e-3,
     learned_noise_scale = True,
-    noise_scale_learning_rate = 1e-5
+    noise_scale_learning_rate = 2e-5
 )
 
 evo_strat()

{x_evolution-0.1.11 → x_evolution-0.1.15}/train_xor.py

@@ -1,6 +1,7 @@
 import torch
 from torch import tensor
 import torch.nn.functional as F
+from torch.optim.lr_scheduler import LambdaLR
 
 # model
 
@@ -44,7 +45,13 @@ evo_strat = EvoStrategy(
     noise_scale = 1e-1,
     noise_scale_clamp_range = (5e-2, 2e-1),
     learned_noise_scale = True,
-    noise_scale_learning_rate = 5e-4
+    noise_scale_learning_rate = 5e-4,
+    use_scheduler = True,
+    scheduler_klass = LambdaLR,
+    scheduler_kwargs = dict(lr_lambda = lambda step: min(1., step / 10.)),
+    use_sigma_scheduler = True,
+    sigma_scheduler_klass = LambdaLR,
+    sigma_scheduler_kwargs = dict(lr_lambda = lambda step: min(1., step / 10.))
 )
 
 evo_strat()

{x_evolution-0.1.11 → x_evolution-0.1.15}/x_evolution/x_evolution.py

@@ -8,7 +8,8 @@ from functools import partial
 import torch
 from torch import tensor, Tensor, is_tensor, arange, randint
 from torch.nn import Module, ModuleList, Parameter, ParameterList
-from torch.optim import SGD, Adam
+from torch.optim import SGD, Adam, Optimizer
+from torch.optim.lr_scheduler import LRScheduler
 
 import torch.nn.functional as F
 
@@ -73,11 +74,17 @@ class EvoStrategy(Module):
         noise_scale_learning_rate = 1e-5,
         noise_scale_clamp_range: tuple[float, float] = (1e-3, 1e-1),
         use_optimizer = True,
-        optimizer_klass = partial(SGD, nesterov = True, momentum = 0.1, weight_decay = 1e-2),
+        optimizer_klass: type[Optimizer] | Callable = partial(SGD, nesterov = True, momentum = 0.1, weight_decay = 1e-2),
         optimizer_kwargs: dict = dict(),
         use_sigma_optimizer = True,
-        sigma_optimizer_klass = partial(SGD, nesterov = True, momentum = 0.1),
+        sigma_optimizer_klass: type[Optimizer] | Callable = partial(SGD, nesterov = True, momentum = 0.1),
         sigma_optimizer_kwargs: dict = dict(),
+        use_scheduler = False,
+        scheduler_klass: type[LRScheduler] | None = None,
+        scheduler_kwargs: dict = dict(),
+        use_sigma_scheduler = False,
+        sigma_scheduler_klass: type[LRScheduler] | None = None,
+        sigma_scheduler_kwargs: dict = dict(),
         transform_fitness: Callable = identity,
         fitness_to_weighted_factor: Callable[[Tensor], Tensor] = normalize,
         checkpoint_every = None,            # saving every number of generations
@@ -198,6 +205,16 @@ class EvoStrategy(Module):
 
         # rejecting the fitnesses for a certain generation if this function is true
 
+        self.use_scheduler = use_scheduler
+
+        if use_scheduler and exists(scheduler_klass) and use_optimizer:
+            self.scheduler = scheduler_klass(self.optimizer, **scheduler_kwargs)
+
+        self.use_sigma_scheduler = use_sigma_scheduler
+
+        if use_sigma_scheduler and exists(sigma_scheduler_klass) and use_sigma_optimizer:
+            self.sigma_scheduler = sigma_scheduler_klass(self.sigma_optimizer, **sigma_scheduler_kwargs)
+
         self.reject_generation_fitnesses_if = reject_generation_fitnesses_if
 
         # verbose
@@ -310,9 +327,6 @@ class EvoStrategy(Module):
 
                     if self.use_sigma_optimizer:
                         accum_grad_(sigma, -one_grad_sigma)
-
-                        self.sigma_optimizer.step()
-                        self.sigma_optimizer.zero_grad()
                     else:
                         sigma.add_(one_grad_sigma * self.noise_scale_learning_rate)
 
@@ -324,11 +338,17 @@ class EvoStrategy(Module):
             self.optimizer.step()
             self.optimizer.zero_grad()
 
+            if self.use_scheduler and exists(self.scheduler):
+                self.scheduler.step()
+
         if self.learned_noise_scale:
             if self.use_sigma_optimizer:
                 self.sigma_optimizer.step()
                 self.sigma_optimizer.zero_grad()
 
+                if self.use_sigma_scheduler and exists(self.sigma_scheduler):
+                    self.sigma_scheduler.step()
+
             for sigma in self.sigmas:
                 self.sigma_clamp_(sigma)
 
@@ -495,7 +515,7 @@ class EvoStrategy(Module):
             if self.learned_noise_scale:
                 packed_sigma, _ = pack(list(self.sigmas), '*')
                 avg_sigma = packed_sigma.mean().item()
-                msg += f' | avg sigma: {avg_sigma:.3f}'
+                msg += f' | average sigma: {avg_sigma:.3f}'
 
             self.print(msg)