evolutionary-policy-optimization 0.2.5__tar.gz → 0.2.6__tar.gz

{evolutionary_policy_optimization-0.2.5 → evolutionary_policy_optimization-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.2.5
+Version: 0.2.6
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization

{evolutionary_policy_optimization-0.2.5 → evolutionary_policy_optimization-0.2.6}/evolutionary_policy_optimization/experimental.py

@@ -1,4 +1,5 @@
 from random import uniform
+from copy import deepcopy
 
 import torch
 import torch.nn.functional as F
@@ -100,6 +101,7 @@ class PopulationWrapper(Module):
         assert num_selected < pop_size
         assert tournament_size < num_selected
 
+        self.pop_size = pop_size
         self.num_selected = num_selected
         self.tournament_size = tournament_size
         self.num_offsprings = pop_size - num_selected
@@ -123,6 +125,14 @@ class PopulationWrapper(Module):
     def pop_params(self):
         return dict(zip(self.param_names, self.param_values))
 
+    def individual(self, id) -> Module:
+        assert 0 <= id < self.pop_size
+        state_dict = {key: param[id] for key, param in self.pop_params.items()}
+
+        net = deepcopy(self.net)
+        net.load_state_dict(state_dict)
+        return net
+
     def parameters(self):
         return self.pop_params.values()
 
@@ -162,9 +172,18 @@ class PopulationWrapper(Module):
     def forward(
         self,
         data,
+        *,
+        individual_id = None,
         labels = None,
         return_logits_with_loss = False
     ):
+        # if `individual_id` passed in, will forward for only that one network
+
+        if exists(individual_id):
+            assert 0 <= individual_id < self.pop_size
+            params = {key: param[individual_id] for key, param in self.pop_params.items()}
+            return functional_call(self.net, params, data)
+
         out = self.forward_pop_nets(dict(self.pop_params), data)
 
         if not exists(labels):

{evolutionary_policy_optimization-0.2.5 → evolutionary_policy_optimization-0.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.2.5"
+version = "0.2.6"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{evolutionary_policy_optimization-0.2.5 → evolutionary_policy_optimization-0.2.6}/train_crossover_weight_space.py

@@ -109,7 +109,7 @@ for i in range(1000):
 
     data, labels = next(iter_train_dl)
 
-    losses = pop_net(data, labels)
+    losses = pop_net(data, labels = labels)
 
     losses.sum(dim = 0).mean().backward()
 
@@ -127,7 +127,7 @@ for i in range(1000):
             pop_net.eval()
 
             eval_data, labels = next(iter_eval_dl)
-            eval_loss, logits = pop_net(eval_data, labels, return_logits_with_loss = True)
+            eval_loss, logits = pop_net(eval_data, labels = labels, return_logits_with_loss = True)
 
             total = labels.shape[0] * pop_size
             correct = (logits.argmax(dim = -1) == labels).long().sum().item()