Compile optimizer after crossover

Author: TimeDelta

genome.py

@@ -4,6 +4,8 @@
 from random import choice, random, shuffle
 from typing import Dict, List, Tuple
 
+import torch
+
 from neat.aggregations import AggregationFunctionSet
 from neat.config import ConfigParameter, write_pretty_params
 from neat.graphs import creates_cycle, required_for_output
@@ -191,6 +193,8 @@ def configure_crossover(self, genome1, genome2, config):
                 # Homologous gene: combine genes from both parents.
                 self.nodes[key] = ng1.crossover(ng2)
 
+
+
     def mutate(self, config):
         """Mutates this genome."""
 
@@ -409,6 +413,19 @@ def get_pruned_copy(self, genome_config):
         new_genome.connections = used_connection_genes
         return new_genome
 
+    def compile_optimizer(self, genome_config):
+        """Compile this genome into a TorchScript optimizer."""
+        from graph_builder import DynamicOptimizerModule
+
+        if not self.connections:
+            self.optimizer = None
+        else:
+            module = DynamicOptimizerModule(
+                self, genome_config.input_keys, genome_config.output_keys, self.graph_dict
+            )
+            self.optimizer = torch.jit.script(module)
+        self.optimizer_path = None
+
     def add_node(self, node_type: str, activation, aggregation) -> NodeGene:
         if activation is None and aggregation is None:
             print("WARNING: node added without any operation")

reproduction.py

@@ -103,6 +103,8 @@ def reproduce(self, config, species, pop_size, generation, task):
                     child = config.genome_type(cid)
                     child.configure_crossover(p1, p2, config.genome_config)
                     child.mutate(config.genome_config)
+                    if hasattr(child, "compile_optimizer"):
+                        child.compile_optimizer(config.genome_config)
                     new_population[cid] = child
                     self.ancestors[cid] = (p1_id, p2_id)