x-evolution 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

x_evolution/x_evolution.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from functools import partial
 
 import torch
-from torch import tensor, Tensor, is_tensor, arange, randint
+from torch import tensor, Tensor, stack, is_tensor, arange, randint
 from torch.nn import Module, ModuleList, Parameter, ParameterList
 from torch.optim import SGD, Adam, Optimizer
 from torch.optim.lr_scheduler import LRScheduler
@@ -366,11 +366,19 @@ class EvoStrategy(Module):
         self,
         filename = 'evolved.model',
         num_generations = None,
-        disable_distributed = False
+        disable_distributed = False,
+        rollback_model_at_end = False,
+        verbose = None
     ):
+        verbose = default(verbose, self.verbose)
 
         model = self.noisable_model.to(self.device)
 
+        # maybe save model for rolling back (for meta-evo)
+
+        if rollback_model_at_end:
+            self.checkpoint('initial.model')
+
         # maybe sigmas
 
         if self.learned_noise_scale:
@@ -533,9 +541,18 @@ class EvoStrategy(Module):
 
         self.print('evolution complete')
 
+        # final checkpoint
+
         self.checkpoint(f'{filename}.final.{generation}')
 
+        # maybe rollback
+
+        if rollback_model_at_end:
+            orig_state_dict = torch.load(str(self.checkpoint_folder / 'initial.model.pt'), weights_only = True)
+
+            self.model.load_state_dict(orig_state_dict)
+
         # return fitnesses across generations
         # for meta-evolutionary (nesting EvoStrategy within the environment of another and optimizing some meta-network)
 
-        return fitnesses_across_generations
+        return stack(fitnesses_across_generations)

{x_evolution-0.1.15.dist-info → x_evolution-0.1.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.15
+Version: 0.1.17
 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.17.dist-info/RECORD

@@ -0,0 +1,6 @@
+x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
+x_evolution/x_evolution.py,sha256=UaDdGLvd2xURQp2SnBjLZkdyOPQ41Y55djo00jnKCNE,18503
+x_evolution-0.1.17.dist-info/METADATA,sha256=KSqDsQ5whxAfgPE6is0xQZTDPAPMJFKSCn6lSGd1Ork,5649
+x_evolution-0.1.17.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+x_evolution-0.1.17.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+x_evolution-0.1.17.dist-info/RECORD,,

x_evolution-0.1.15.dist-info/RECORD

@@ -1,6 +0,0 @@
-x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
-x_evolution/x_evolution.py,sha256=zW2sSev-jDtkOFNScx4p2Ncrdqj39SYef2ynApjo8RU,17979
-x_evolution-0.1.15.dist-info/METADATA,sha256=apdgiJ4-mMt2MU034eLxu6SJkQ-eeRJJqx7w7G9l0YI,5649
-x_evolution-0.1.15.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-x_evolution-0.1.15.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-x_evolution-0.1.15.dist-info/RECORD,,