Changed checkpoint dump from end_generation to post_evaluate; added best_genome to checkpoint and Population initial_state

hugoaboud · hugoaboud · commit e70a0689ede2 · 2020-12-28T00:09:12.000-03:00
diff --git a/neat/checkpoint.py b/neat/checkpoint.py
@@ -43,7 +43,7 @@ def __init__(self, generation_interval=100, time_interval_seconds=300,
     def start_generation(self, generation):
         self.current_generation = generation
 
-    def end_generation(self, config, population, species_set):
+    def post_evaluate(self, config, population, species_set, best_genome):
         checkpoint_due = False
 
         if self.time_interval_seconds is not None:
@@ -57,23 +57,23 @@ def end_generation(self, config, population, species_set):
                 checkpoint_due = True
 
         if checkpoint_due:
-            self.save_checkpoint(config, population, species_set, self.current_generation)
+            self.save_checkpoint(config, population, species_set, best_genome, self.current_generation)
             self.last_generation_checkpoint = self.current_generation
             self.last_time_checkpoint = time.time()
 
-    def save_checkpoint(self, config, population, species_set, generation):
+    def save_checkpoint(self, config, population, species_set, best_genome, generation):
         """ Save the current simulation state. """
         filename = '{0}{1}'.format(self.filename_prefix, generation)
         print("Saving checkpoint to {0}".format(filename))
 
         with gzip.open(filename, 'w', compresslevel=5) as f:
-            data = (generation, config, population, species_set, random.getstate())
+            data = (generation, config, population, species_set, best_genome, random.getstate())
             pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
 
     @staticmethod
     def restore_checkpoint(filename):
         """Resumes the simulation from a previous saved point."""
         with gzip.open(filename) as f:
-            generation, config, population, species_set, rndstate = pickle.load(f)
+            generation, config, population, species_set, best_genome, rndstate = pickle.load(f)
             random.setstate(rndstate)
-            return Population(config, (population, species_set, generation))
+            return Population(config, (population, species_set, best_genome, generation))
diff --git a/neat/population.py b/neat/population.py
@@ -44,10 +44,10 @@ def __init__(self, config, initial_state=None):
             self.species = config.species_set_type(config.species_set_config, self.reporters)
             self.generation = 0
             self.species.speciate(config, self.population, self.generation)
+            self.best_genome = None
         else:
-            self.population, self.species, self.generation = initial_state
+            self.population, self.species, self.best_genome, self.generation = initial_state
 
-        self.best_genome = None
 
     def add_reporter(self, reporter):
         self.reporters.add(reporter)
