evolutionary-policy-optimization 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -382,6 +382,52 @@ class StateNorm(Module):
 
         return normed
 
+# style mapping network from StyleGAN2
+
+class EqualLinear(Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        lr_mul = 1,
+        bias = True
+    ):
+        super().__init__()
+        self.lr_mul = lr_mul
+
+        self.weight = nn.Parameter(torch.randn(dim_out, dim_in))
+        self.bias = nn.Parameter(torch.zeros(dim_out))
+
+    def forward(
+        self,
+        input
+    ):
+        weight, bias = tuple(t * self.lr_mul for t in (self.weight, self.bias))
+        return F.linear(input, weight, bias = bias)
+
+class LatentMappingNetwork(Module):
+    def __init__(
+        self,
+        dim_latent,
+        depth,
+        lr_mul = 0.1,
+        leaky_relu_p = 2e-2
+    ):
+        super().__init__()
+
+        layers = []
+
+        for i in range(depth):
+            layers.extend([
+                EqualLinear(dim_latent, dim_latent, lr_mul),
+                nn.LeakyReLU(leaky_relu_p)
+            ])
+
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.net(x)
+
 # simple MLP networks, but with latent variables
 # the latent variables are the "genes" with the rest of the network as the scaffold for "gene expression" - as suggested in the paper
 
@@ -391,6 +437,7 @@ class MLP(Module):
         dim,
         depth,
         dim_latent = 0,
+        latent_mapping_network_depth = 2,
         expansion_factor = 2.
     ):
         super().__init__()
@@ -401,6 +448,7 @@ class MLP(Module):
         self.needs_latent = dim_latent > 0
 
         self.encode_latent = nn.Sequential(
+            LatentMappingNetwork(dim_latent, depth = latent_mapping_network_depth),
             Linear(dim_latent, dim * 2),
             nn.SiLU()
         ) if self.needs_latent else None

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.dist-info → evolutionary_policy_optimization-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.2.5
+Version: 0.2.7
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -268,4 +268,14 @@ That's it
 }
 ```
 
+```bibtex
+@article{Karras2019stylegan2,
+    title   = {Analyzing and Improving the Image Quality of {StyleGAN}},
+    author  = {Tero Karras and Samuli Laine and Miika Aittala and Janne Hellsten and Jaakko Lehtinen and Timo Aila},
+    journal = {CoRR},
+    volume  = {abs/1912.04958},
+    year    = {2019},
+}
+```
+
 *Evolution is cleverer than you are.* - Leslie Orgel

{evolutionary_policy_optimization-0.2.5.dist-info → evolutionary_policy_optimization-0.2.7.dist-info}/RECORD

@@ -1,10 +1,10 @@
 evolutionary_policy_optimization/__init__.py,sha256=NyiYDYU7DlpmOTM7xiBQET3r1WwX0ebrgMCBLSQrW3c,288
 evolutionary_policy_optimization/distributed.py,sha256=MxyxqxANAuOm8GYb0Yu09EHd_aVLhK2uwgrfuVWciPU,2342
 evolutionary_policy_optimization/env_wrappers.py,sha256=bDL06o9_b1iW6k3fw2xifnOnYlzs643tdW6Yv2gsIdw,803
-evolutionary_policy_optimization/epo.py,sha256=81rf249ykaPrAEMGk9KsF98qDkCUhW8xL3-2UXIvI2E,51838
-evolutionary_policy_optimization/experimental.py,sha256=NKtAeOS8AVoX5HHwwxAl1ngUi7ZE19uV1-NILFK6Tu8,4877
+evolutionary_policy_optimization/epo.py,sha256=OKumVrOH7DSKZhbbx-5oCI_JJwJNYf4lrEpCNmbj6ZY,52991
+evolutionary_policy_optimization/experimental.py,sha256=7LOrMIaU4fr2Vme1ZpHNIvlvFEIdWj0-uemhQoNJcPQ,5549
 evolutionary_policy_optimization/mock_env.py,sha256=TLyyRm6tOD0Kdn9QqJJQriaSnsR-YmNQHo4OohmZFG4,1410
-evolutionary_policy_optimization-0.2.5.dist-info/METADATA,sha256=WzU43HLqNw6kjQ-s7IMW9bFv9nojiQsLEd2pHTvKbfw,8858
-evolutionary_policy_optimization-0.2.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.2.5.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.2.5.dist-info/RECORD,,
+evolutionary_policy_optimization-0.2.7.dist-info/METADATA,sha256=MiG_AYp6KoANhdGuaGM37-zaciW8dCrO0KvkXk7hO7w,9171
+evolutionary_policy_optimization-0.2.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.2.7.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.2.7.dist-info/RECORD,,