statedict2pytree 0.3.0__tar.gz → 0.4.0__tar.gz

{statedict2pytree-0.3.0 → statedict2pytree-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: statedict2pytree
-Version: 0.3.0
+Version: 0.4.0
 Summary: Converts torch models into PyTrees for Equinox
 Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
 Requires-Python: ~=3.10

statedict2pytree-0.4.0/examples/convert_resnet.py

@@ -0,0 +1,20 @@
+import equinox as eqx
+import jax
+import statedict2pytree as s2p
+from resnet import resnet152
+from torchvision.models import resnet152 as t_resnet152, ResNet152_Weights
+
+
+def convert_resnet():
+    resnet_jax = resnet152(key=jax.random.PRNGKey(33), make_with_state=False)
+    resnet_torch = t_resnet152(weights=ResNet152_Weights.DEFAULT)
+    state_dict = resnet_torch.state_dict()
+
+    # s2p.start_conversion(resnet_jax, state_dict)
+    model, state = s2p.autoconvert(resnet_jax, state_dict)
+    name = "resnet152.eqx"
+    eqx.tree_serialise_leaves(name, (model, state))
+
+
+if __name__ == "__main__":
+    convert_resnet()

{statedict2pytree-0.3.0 → statedict2pytree-0.4.0}/examples/resnet.py

@@ -333,12 +333,22 @@ def resnet18(
 
 
 def resnet34(
-    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+    image_channels: int = 3,
+    num_classes: int = 1000,
+    *,
+    key: PRNGKeyArray,
+    make_with_state: bool = True,
+    **kwargs,
 ):
     layers = [3, 4, 6, 3]
-    return eqx.nn.make_with_state(ResNet)(
-        BasicBlock, layers, image_channels, num_classes, **kwargs, key=key
-    )
+    if make_with_state:
+        return eqx.nn.make_with_state(ResNet)(
+            BasicBlock, layers, image_channels, num_classes, **kwargs, key=key
+        )
+    else:
+        return ResNet(
+            BasicBlock, layers, image_channels, num_classes, **kwargs, key=key
+        )
 
 
 def resnet50(
@@ -361,18 +371,39 @@ def resnet50(
 
 
 def resnet101(
-    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+    image_channels: int = 3,
+    num_classes: int = 1000,
+    *,
+    key: PRNGKeyArray,
+    make_with_state: bool = True,
+    **kwargs,
 ):
     layers = [3, 4, 23, 3]
-    return eqx.nn.make_with_state(ResNet)(
-        Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
-    )
+
+    if make_with_state:
+        return eqx.nn.make_with_state(ResNet)(
+            Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+        )
+    else:
+        return ResNet(
+            Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+        )
 
 
 def resnet152(
-    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+    image_channels: int = 3,
+    num_classes: int = 1000,
+    *,
+    key: PRNGKeyArray,
+    make_with_state: bool = True,
+    **kwargs,
 ):
     layers = [3, 8, 36, 3]
-    return eqx.nn.make_with_state(ResNet)(
-        Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
-    )
+    if make_with_state:
+        return eqx.nn.make_with_state(ResNet)(
+            Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+        )
+    else:
+        return ResNet(
+            Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+        )

{statedict2pytree-0.3.0 → statedict2pytree-0.4.0}/examples/test_resnet_inference.py

@@ -6,15 +6,15 @@ import equinox as eqx
 import jax
 import jax.numpy as jnp
 import torch
+from examples.resnet import resnet152
 from PIL import Image
-from tests.resnet import resnet50
 from torchvision import transforms
-from torchvision.models import resnet50 as t_resnet50, ResNet50_Weights
+from torchvision.models import resnet152 as t_resnet152, ResNet152_Weights
 
 
 def test_resnet():
-    resnet_jax = resnet50(key=jax.random.PRNGKey(33), make_with_state=False)
-    resnet_torch = t_resnet50(weights=ResNet50_Weights.DEFAULT)
+    resnet_jax = resnet152(key=jax.random.PRNGKey(33), make_with_state=False)
+    resnet_torch = t_resnet152(weights=ResNet152_Weights.DEFAULT)
 
     img_name = "doggo.jpeg"
 
@@ -42,7 +42,7 @@ def test_resnet():
     )  # Outputs the ImageNet class index of the prediction
 
     url = "https://storage.googleapis.com/download.tensorflow.org/data/imagenet_class_index.json"
-    with urllib.request.urlopen(url) as url:
+    with urllib.request.urlopen(url) as url:  # pyright: ignore
         imagenet_labels = json.loads(url.read().decode())
 
     label = imagenet_labels[str(predicted.item())][1]
@@ -52,7 +52,7 @@ def test_resnet():
     model_callable = ft.partial(identity, resnet_jax)
     model, state = eqx.nn.make_with_state(model_callable)()
 
-    model, state = eqx.tree_deserialise_leaves("model.eqx", (model, state))
+    model, state = eqx.tree_deserialise_leaves("resnet152.eqx", (model, state))
 
     jax_batch = jnp.array(batch_t.numpy())
     out, state = eqx.filter_vmap(

{statedict2pytree-0.3.0 → statedict2pytree-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "statedict2pytree"
-version = "0.3.0"
+version = "0.4.0"
 description = "Converts torch models into PyTrees for Equinox"
 readme = "README.md"
 requires-python = "~=3.10"

{statedict2pytree-0.3.0 → statedict2pytree-0.4.0}/statedict2pytree/statedict2pytree.py

@@ -46,7 +46,7 @@ def get_node(
         return tree
     else:
         next_target: str = targets[0]
-        if bool(re.search(r"\[\d\]", next_target)):
+        if bool(re.search(r"\[\d+\]", next_target)):
             split_index = next_target.rfind("[")
             name, index = next_target[:split_index], next_target[split_index:]
             index = index[1:-1]
@@ -157,6 +157,9 @@ def index():
 def autoconvert(pytree: PyTree, state_dict: dict) -> tuple[PyTree, eqx.nn.State]:
     jax_fields = pytree_to_fields(pytree)
     torch_fields = state_dict_to_fields(state_dict)
+
+    for k, v in state_dict.items():
+        state_dict[k] = v.numpy()
     return convert(jax_fields, torch_fields, pytree, state_dict)
 
 

statedict2pytree-0.3.0/examples/convert_resnet.py

@@ -1,16 +0,0 @@
-import jax
-import statedict2pytree as s2p
-from resnet import resnet50
-from torchvision.models import resnet50 as t_resnet50, ResNet50_Weights
-
-
-def convert_resnet():
-    resnet_jax = resnet50(key=jax.random.PRNGKey(33), make_with_state=False)
-    resnet_torch = t_resnet50(weights=ResNet50_Weights.DEFAULT)
-    state_dict = resnet_torch.state_dict()
-
-    s2p.start_conversion(resnet_jax, state_dict)
-
-
-if __name__ == "__main__":
-    convert_resnet()