Merge pull request #44 from RobertTLange/develop

Fix checkpoint loading for LES

CHANGELOG.md

@@ -11,6 +11,10 @@
     - [ ] Wavelet Based Encoding (van Steenkiste, 2016)
     - [ ] CNN Hypernetwork (Ha - start with simple MLP)
 
+### [v0.1.2] - [03/2023]
+
+- Fix LES checkpoint loading from package data via `pkgutil`.
+
 ### [v0.1.1] - [03/2023]
 
 ##### Added

evosax/_version.py

@@ -1 +1 @@
-__version__ = "0.1.1"
+__version__ = "0.1.2"

evosax/strategies/les.py

@@ -4,6 +4,7 @@
 from flax import struct
 import jax
 import jax.numpy as jnp
+import pkgutil
 from ..utils.learned_eo import (
     AttentionWeights,
     EvoPathMLP,
@@ -77,7 +78,8 @@ def __init__(
             net_ckpt_path = os.path.join(
                 os.path.dirname(__file__), f"ckpt/{ckpt_fname}"
             )
-            self.les_net_params = load_pkl_object(net_ckpt_path)
+            data = pkgutil.get_data(__name__, f"ckpt/{ckpt_fname}")
+            self.les_net_params = load_pkl_object(data, pkg_load=True)
             print(f"Loaded pretrained LES model from ckpt: {ckpt_fname}")
 
     @property

evosax/utils/learned_eo.py

@@ -1,3 +1,4 @@
+from typing import Any
 import functools
 import sys
 import chex
@@ -18,10 +19,13 @@
     import pickle
 
 
-def load_pkl_object(filename: str) -> chex.ArrayTree:
+def load_pkl_object(filename: Any, pkg_load: bool = False) -> chex.ArrayTree:
     """Reload pickle objects from path."""
-    with open(filename, "rb") as input:
-        obj = pickle.load(input)
+    if not pkg_load:
+        with open(filename, "rb") as input:
+            obj = pickle.load(input)
+    else:
+        obj = pickle.loads(filename)
     return obj