statedict2pytree 0.2.0__tar.gz → 0.4.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: statedict2pytree
-Version: 0.2.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
@@ -32,6 +32,12 @@ Usually, if you _declared the fields in the same order as in the PyTorch model_,
 
 (Theoretically, you can rearrange the model in any way you like - e.g. last layer as the first layer - as long as the shapes match!)
 
+## Shape Matching? What's that?
+
+Currently, there is no sophisticated shape matching in place. Two matrices are considered "matching" if the product of their shape match. For example:
+
+1. (8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
+
 ## Get Started
 
 ### Installation
@@ -124,3 +130,18 @@ def convert(
 ```
 
 If your models already have the right "order", then you might as well use this function directly. Note that the lists `jax_fields` and `torch_fields` must have the same length and each matching entry must have the same shape!
+
+For the full, automatic experience, use `autoconvert`:
+
+```python
+import statedict2pytree as s2p
+
+my_model = Model(...)
+state_dict = ...
+
+model, state = s2p.autoconvert(my_model, state_dict)
+
+```
+
+This will however only work if your PyTree fields have been declared
+in the same order as they appear in the state dict!

{statedict2pytree-0.2.0 → statedict2pytree-0.4.0}/README.md

@@ -8,6 +8,12 @@ Usually, if you _declared the fields in the same order as in the PyTorch model_,
 
 (Theoretically, you can rearrange the model in any way you like - e.g. last layer as the first layer - as long as the shapes match!)
 
+## Shape Matching? What's that?
+
+Currently, there is no sophisticated shape matching in place. Two matrices are considered "matching" if the product of their shape match. For example:
+
+1. (8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
+
 ## Get Started
 
 ### Installation
@@ -100,3 +106,18 @@ def convert(
 ```
 
 If your models already have the right "order", then you might as well use this function directly. Note that the lists `jax_fields` and `torch_fields` must have the same length and each matching entry must have the same shape!
+
+For the full, automatic experience, use `autoconvert`:
+
+```python
+import statedict2pytree as s2p
+
+my_model = Model(...)
+state_dict = ...
+
+model, state = s2p.autoconvert(my_model, state_dict)
+
+```
+
+This will however only work if your PyTree fields have been declared
+in the same order as they appear in the state dict!

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.2.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.2.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.2.0 → statedict2pytree-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "statedict2pytree"
-version = "0.2.0"
+version = "0.4.0"
 description = "Converts torch models into PyTrees for Equinox"
 readme = "README.md"
 requires-python = "~=3.10"
@@ -16,7 +16,7 @@ dependencies = [
     "torch",
     "flask",
     "pydantic",
-    "penzai"
+    "penzai",
 ]
 [project.optional-dependencies]
 dev = ["nox", "pre-commit", "pytest", "mkdocs"]

statedict2pytree-0.4.0/pyrightconfig.json

@@ -0,0 +1,4 @@
+{
+  "venvPath": ".",
+  "venv": ".venv"
+}

statedict2pytree-0.4.0/statedict2pytree/__init__.py

@@ -0,0 +1,7 @@
+from statedict2pytree.statedict2pytree import (
+    autoconvert as autoconvert,
+    convert as convert,
+    pytree_to_fields as pytree_to_fields,
+    start_conversion as start_conversion,
+    state_dict_to_fields as state_dict_to_fields,
+)

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

@@ -4,6 +4,7 @@ import re
 import equinox as eqx
 import flask
 import jax
+import numpy as np
 from beartype.typing import Optional
 from jaxtyping import PyTree
 from loguru import logger
@@ -31,6 +32,13 @@ PYTREE: Optional[PyTree] = None
 STATE_DICT: Optional[dict] = None
 
 
+def can_reshape(shape1, shape2):
+    product1 = np.prod(shape1)
+    product2 = np.prod(shape2)
+
+    return product1 == product2
+
+
 def get_node(
     tree: PyTree, targets: list[str], log_when_not_found: bool = False
 ) -> PyTree | None:
@@ -38,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]
@@ -146,17 +154,26 @@ 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)
+
+
 def convert(
     jax_fields: list[JaxField],
     torch_fields: list[TorchField],
     pytree: PyTree,
     state_dict: dict,
-):
+) -> tuple[PyTree, eqx.nn.State]:
     identity = lambda *args, **kwargs: pytree
     model, state = eqx.nn.make_with_state(identity)()
     state_paths: list[tuple[JaxField, TorchField]] = []
     for jax_field, torch_field in zip(jax_fields, torch_fields):
-        if jax_field.shape != torch_field.shape:
+        if not can_reshape(jax_field.shape, torch_field.shape):
             raise ValueError(
                 "Fields have incompatible shapes!"
                 f"{jax_field.shape=} != {torch_field.shape=}"
@@ -171,7 +188,7 @@ def convert(
                 model = eqx.tree_at(
                     where,
                     model,
-                    state_dict[torch_field.path],
+                    state_dict[torch_field.path].reshape(jax_field.shape),
                 )
     result: dict[str, list[TorchField]] = {}
     for tuple_item in state_paths:

statedict2pytree-0.4.0/tests/test_conv.py

@@ -0,0 +1,61 @@
+import equinox as eqx
+import jax
+import numpy as np
+import statedict2pytree as s2p
+import torch
+
+
+def test_conv():
+    in_channels = 8
+    out_channels = 8
+    kernel_size = 4
+    stride = 2
+    padding = 1
+
+    class J(eqx.Module):
+        conv: eqx.nn.Conv2d
+
+        def __init__(self):
+            self.conv = eqx.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                key=jax.random.PRNGKey(22),
+            )
+
+    class T(torch.nn.Module):
+        def __init__(self) -> None:
+            super(T, self).__init__()
+            self.conv = torch.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+            )
+
+    jax_model = J()
+    torch_model = T()
+    state_dict = torch_model.state_dict()
+
+    jax_fields = s2p.pytree_to_fields(jax_model)
+    torch_fields = s2p.state_dict_to_fields(state_dict)
+
+    model, state = s2p.convert(
+        jax_fields, torch_fields, pytree=jax_model, state_dict=state_dict
+    )
+
+    assert np.allclose(
+        np.array(model.conv.weight), torch_model.conv.weight.detach().numpy()
+    )
+    if torch_model.conv.bias is not None:
+        assert np.allclose(
+            np.array(model.conv.bias),
+            torch_model.conv.bias.detach().numpy().reshape(model.conv.bias.shape),
+        )
+
+
+if __name__ == "__main__":
+    test_conv()

statedict2pytree-0.4.0/tests/test_linear.py

@@ -0,0 +1,45 @@
+import equinox as eqx
+import jax
+import numpy as np
+import statedict2pytree as s2p
+import torch
+
+
+def test_linear():
+    in_features = 10
+    out_features = 10
+
+    class J(eqx.Module):
+        linear: eqx.nn.Linear
+
+        def __init__(self):
+            self.linear = eqx.nn.Linear(
+                in_features, out_features, key=jax.random.PRNGKey(30)
+            )
+
+    class T(torch.nn.Module):
+        def __init__(self) -> None:
+            super(T, self).__init__()
+            self.linear = torch.nn.Linear(in_features, out_features)
+
+    jax_model = J()
+    torch_model = T()
+    state_dict = torch_model.state_dict()
+
+    jax_fields = s2p.pytree_to_fields(jax_model)
+    torch_fields = s2p.state_dict_to_fields(state_dict)
+
+    model, state = s2p.convert(
+        jax_fields, torch_fields, pytree=jax_model, state_dict=state_dict
+    )
+
+    assert np.allclose(
+        np.array(model.linear.weight), torch_model.linear.weight.detach().numpy()
+    )
+    assert np.allclose(
+        np.array(model.linear.bias), torch_model.linear.bias.detach().numpy()
+    )
+
+
+if __name__ == "__main__":
+    test_linear()

statedict2pytree-0.2.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()

statedict2pytree-0.2.0/statedict2pytree/__init__.py

@@ -1,4 +0,0 @@
-from statedict2pytree.statedict2pytree import (
-    convert as convert,
-    start_conversion as start_conversion,
-)