statedict2pytree 0.1.2__tar.gz

statedict2pytree-0.1.2/.gitignore

@@ -0,0 +1,161 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+node_modules/

statedict2pytree-0.1.2/.pre-commit-config.yaml

@@ -0,0 +1,16 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.4.4
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args: [--fix]
+      # Run the formatter.
+      - id: ruff-format
+  - repo: https://github.com/RobertCraigie/pyright-python
+    rev: v1.1.351
+    hooks:
+      - id: pyright
+        additional_dependencies:
+          [beartype, jax, jaxtyping, pytest, typing_extensions]

statedict2pytree-0.1.2/PKG-INFO

@@ -0,0 +1,43 @@
+Metadata-Version: 2.3
+Name: statedict2pytree
+Version: 0.1.2
+Summary: Converts torch models into PyTrees for Equinox
+Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
+Requires-Python: ~=3.10
+Requires-Dist: beartype
+Requires-Dist: equinox>=0.11.4
+Requires-Dist: flask
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: jaxtyping
+Requires-Dist: loguru
+Requires-Dist: pydantic
+Requires-Dist: torch
+Requires-Dist: typing-extensions
+Provides-Extra: dev
+Requires-Dist: mkdocs; extra == 'dev'
+Requires-Dist: nox; extra == 'dev'
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# statedict2pytree
+
+![statedict2pytree](torch2jax.png "A ResNet demo")
+
+The goal of this package is to simplify the conversion from PyTorch models into JAX PyTrees (which can be used e.g. in Equinox). The way this works is by putting both models side my side and aligning the weights in the right order. Then, all statedict2pytree is doing, is iterating over both lists and matching the weight matrices.
+
+Usually, if you _declared the fields in the same order as in the PyTorch model_, you don't have to rearrange anything -- but the option is there if you need it.
+
+(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!)
+
+## Get Started
+
+### Installation
+
+Run
+
+```bash
+pip install statedict2pytree
+
+```

statedict2pytree-0.1.2/README.md

@@ -0,0 +1,20 @@
+# statedict2pytree
+
+![statedict2pytree](torch2jax.png "A ResNet demo")
+
+The goal of this package is to simplify the conversion from PyTorch models into JAX PyTrees (which can be used e.g. in Equinox). The way this works is by putting both models side my side and aligning the weights in the right order. Then, all statedict2pytree is doing, is iterating over both lists and matching the weight matrices.
+
+Usually, if you _declared the fields in the same order as in the PyTorch model_, you don't have to rearrange anything -- but the option is there if you need it.
+
+(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!)
+
+## Get Started
+
+### Installation
+
+Run
+
+```bash
+pip install statedict2pytree
+
+```

statedict2pytree-0.1.2/examples/convert_resnet.py

@@ -0,0 +1,16 @@
+import jax
+import statedict2pytree as s2p
+from tests.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.1.2/examples/resnet.py

@@ -0,0 +1,378 @@
+import equinox as eqx
+import jax
+from beartype.typing import Optional, Type
+from equinox.nn import State
+from jaxtyping import Array, PRNGKeyArray
+
+
+def conv3x3(
+    in_channels: int,
+    out_channels: int,
+    stride: int = 1,
+    groups: int = 1,
+    *,
+    key: PRNGKeyArray,
+) -> eqx.nn.Conv2d:
+    return eqx.nn.Conv2d(
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        groups=groups,
+        use_bias=False,
+        key=key,
+    )
+
+
+def conv1x1(
+    in_channels: int, out_channels: int, stride: int = 1, *, key: PRNGKeyArray
+) -> eqx.nn.Conv2d:
+    return eqx.nn.Conv2d(
+        in_channels, out_channels, kernel_size=1, stride=stride, use_bias=False, key=key
+    )
+
+
+class Downsample(eqx.Module):
+    conv: eqx.nn.Conv2d
+    norm: eqx.nn.BatchNorm
+
+    def __init__(
+        self, in_channels: int, out_channels: int, stride: int, *, key: PRNGKeyArray
+    ) -> None:
+        self.conv = conv1x1(in_channels, out_channels, stride, key=key)
+        self.norm = eqx.nn.BatchNorm(out_channels, axis_name="batch")
+
+    def __call__(self, x: Array, state: State) -> tuple[Array, State]:
+        x = self.conv(x)
+        x, state = self.norm(x, state)
+
+        return x, state
+
+
+class BasicBlock(eqx.Module):
+    conv1: eqx.nn.Conv2d
+    bn1: eqx.nn.BatchNorm
+    conv2: eqx.nn.Conv2d
+    bn2: eqx.nn.BatchNorm
+    downsample: Optional[Downsample]
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        stride: int = 1,
+        downsample: Optional[Downsample] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        *,
+        key: PRNGKeyArray,
+    ):
+        if groups != 1 or base_width != 64:
+            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
+        key, conv1_key, conv2_key = jax.random.split(key, 3)
+        self.conv1 = conv3x3(in_channels, out_channels, stride, key=conv1_key)
+        self.bn1 = eqx.nn.BatchNorm(out_channels, axis_name="batch")
+        self.conv2 = conv3x3(out_channels, out_channels, key=conv2_key)
+        self.bn2 = eqx.nn.BatchNorm(out_channels, axis_name="batch")
+        self.downsample = downsample
+
+    def __call__(self, x: Array, state: State) -> tuple[Array, State]:
+        identity = x
+
+        out = self.conv1(x)
+        out, state = self.bn1(out, state)
+        out = jax.nn.relu(out)
+
+        out = self.conv2(out)
+        out, state = self.bn2(out, state)
+
+        if self.downsample is not None:
+            identity, state = self.downsample(x, state)
+
+        out += identity
+        out = jax.nn.relu(out)
+
+        return out, state
+
+
+class Bottleneck(eqx.Module):
+    conv1: eqx.nn.Conv2d
+    bn1: eqx.nn.BatchNorm
+    conv2: eqx.nn.Conv2d
+    bn2: eqx.nn.BatchNorm
+    conv3: eqx.nn.Conv2d
+    bn3: eqx.nn.BatchNorm
+
+    downsample: Optional[Downsample]
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        stride: int = 1,
+        downsample: Optional[Downsample] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        *,
+        key: PRNGKeyArray,
+    ) -> None:
+        width = int(out_channels * (base_width / 64.0)) * groups
+        conv1_key, conv2_key, conv3_key = jax.random.split(key, 3)
+        expansion = 4
+        self.conv1 = conv1x1(in_channels, width, key=conv1_key)
+        self.bn1 = eqx.nn.BatchNorm(width, axis_name="batch")
+        self.conv2 = conv3x3(width, width, stride, groups, key=conv2_key)
+        self.bn2 = eqx.nn.BatchNorm(width, axis_name="batch")
+        self.conv3 = conv1x1(width, out_channels * expansion, key=conv3_key)
+        self.bn3 = eqx.nn.BatchNorm(out_channels * expansion, axis_name="batch")
+        self.downsample = downsample
+
+    def __call__(self, x: Array, state: State) -> tuple[Array, State]:
+        identity = x
+        x = self.conv1(x)
+        x, state = self.bn1(x, state)
+        x = jax.nn.relu(x)
+
+        x = self.conv2(x)
+        x, state = self.bn2(x, state)
+        x = jax.nn.relu(x)
+
+        x = self.conv3(x)
+        x, state = self.bn3(x, state)
+
+        if self.downsample is not None:
+            identity, state = self.downsample(identity, state)
+
+        x += identity
+        x = jax.nn.relu(x)
+
+        return x, state
+
+
+class ResnetLayer(eqx.Module):
+    layers: list[BasicBlock | Bottleneck]
+
+    def __init__(self, layers: list[BasicBlock | Bottleneck]) -> None:
+        self.layers = layers
+
+    def __call__(self, x: Array, state: State) -> tuple[Array, State]:
+        for l in self.layers:
+            x, state = l(x, state)
+        return x, state
+
+
+class ResNet(eqx.Module):
+    in_channels: int = eqx.field(static=True)
+
+    conv1: eqx.nn.Conv2d
+    bn1: eqx.nn.BatchNorm
+    maxpool: eqx.nn.MaxPool2d
+
+    layer1: ResnetLayer
+    layer2: ResnetLayer
+    layer3: ResnetLayer
+    layer4: ResnetLayer
+
+    avgpool: eqx.nn.AdaptiveAvgPool2d
+    fc: eqx.nn.Linear
+
+    def __init__(
+        self,
+        block: Type[BasicBlock | Bottleneck],
+        layers: list[int],
+        image_channels: int = 3,
+        num_classes: int = 1000,
+        groups: int = 1,
+        width_per_group: int = 64,
+        *,
+        key: PRNGKeyArray,
+    ) -> None:
+        self.in_channels = 64
+        key, conv_key = jax.random.split(key)
+        self.conv1 = eqx.nn.Conv2d(
+            image_channels,
+            self.in_channels,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            use_bias=False,
+            key=conv_key,
+        )
+        self.bn1 = eqx.nn.BatchNorm(self.in_channels, axis_name="batch")
+        self.maxpool = eqx.nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        key, *layer_keys = jax.random.split(key, len(layers) + 1)
+        self.layer1 = self._make_layer(
+            block,
+            out_channels=64,
+            num_residual_blocks=layers[0],
+            stride=1,
+            groups=groups,
+            base_width=width_per_group,
+            key=layer_keys[0],
+        )
+        self.layer2 = self._make_layer(
+            block,
+            out_channels=128,
+            num_residual_blocks=layers[1],
+            stride=2,
+            groups=groups,
+            base_width=width_per_group,
+            key=layer_keys[1],
+        )
+
+        self.layer3 = self._make_layer(
+            block,
+            out_channels=256,
+            num_residual_blocks=layers[2],
+            stride=2,
+            groups=groups,
+            base_width=width_per_group,
+            key=layer_keys[2],
+        )
+
+        self.layer4 = self._make_layer(
+            block,
+            out_channels=512,
+            num_residual_blocks=layers[3],
+            stride=2,
+            groups=groups,
+            base_width=width_per_group,
+            key=layer_keys[3],
+        )
+
+        self.avgpool = eqx.nn.AdaptiveAvgPool2d((1, 1))
+        key, fc_key = jax.random.split(key)
+        self.fc = eqx.nn.Linear(512 * _get_expansion(block), num_classes, key=fc_key)
+
+    def __call__(self, x: Array, state: State) -> tuple[Array, State]:
+        x = self.conv1(x)
+        x, state = self.bn1(x, state)
+        x = jax.nn.relu(x)
+        x = self.maxpool(x)
+
+        x, state = self.layer1(x, state)
+        x, state = self.layer2(x, state)
+        x, state = self.layer3(x, state)
+        x, state = self.layer4(x, state)
+        x = self.avgpool(x)
+        x = x.reshape(-1)
+        x = self.fc(x)
+        return x, state
+
+    def _make_layer(
+        self,
+        block: Type[BasicBlock | Bottleneck],
+        out_channels: int,
+        num_residual_blocks: int,
+        stride: int,
+        groups: int,
+        base_width: int,
+        *,
+        key: PRNGKeyArray,
+    ):
+        downsample = None
+        expansion = _get_expansion(block)
+        key, downsample_key = jax.random.split(key)
+        if stride != 1 or self.in_channels != out_channels * expansion:
+            downsample = Downsample(
+                self.in_channels, out_channels * expansion, stride, key=downsample_key
+            )
+        layers = []
+        key, *layer_keys = jax.random.split(key, num_residual_blocks + 1)
+
+        layers.append(
+            block(
+                self.in_channels,
+                out_channels,
+                stride,
+                downsample,
+                groups=groups,
+                base_width=base_width,
+                key=layer_keys[0],
+            )
+        )
+        self.in_channels = out_channels * expansion
+        for i in range(num_residual_blocks - 1):
+            layers.append(
+                block(
+                    self.in_channels,
+                    out_channels,
+                    groups=groups,
+                    base_width=base_width,
+                    key=layer_keys[i + 1],
+                )
+            )
+        return ResnetLayer(layers)
+
+
+def _get_expansion(block_type: Type[Bottleneck | BasicBlock]) -> int:
+    if block_type == Bottleneck:
+        return 4
+    else:
+        return 1
+
+
+def resnet18(
+    image_channels: int = 3,
+    num_classes: int = 1000,
+    *,
+    key: PRNGKeyArray,
+    make_with_state: bool = True,
+    **kwargs,
+):
+    layers = [2, 2, 2, 2]
+    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 resnet34(
+    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+):
+    layers = [3, 4, 6, 3]
+    return eqx.nn.make_with_state(ResNet)(
+        BasicBlock, layers, image_channels, num_classes, **kwargs, key=key
+    )
+
+
+def resnet50(
+    image_channels: int = 3,
+    num_classes: int = 1000,
+    *,
+    key: PRNGKeyArray,
+    make_with_state: bool = True,
+    **kwargs,
+):
+    layers = [3, 4, 6, 3]
+    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 resnet101(
+    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+):
+    layers = [3, 4, 23, 3]
+    return eqx.nn.make_with_state(ResNet)(
+        Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+    )
+
+
+def resnet152(
+    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+):
+    layers = [3, 8, 36, 3]
+    return eqx.nn.make_with_state(ResNet)(
+        Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+    )

statedict2pytree-0.1.2/examples/test_resnet_inference.py

@@ -0,0 +1,68 @@
+import functools as ft
+import json
+import urllib
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import torch
+from PIL import Image
+from tests.resnet import resnet50
+from torchvision import transforms
+from torchvision.models import resnet50 as t_resnet50, ResNet50_Weights
+
+
+def test_resnet():
+    resnet_jax = resnet50(key=jax.random.PRNGKey(33), make_with_state=False)
+    resnet_torch = t_resnet50(weights=ResNet50_Weights.DEFAULT)
+
+    img_name = "doggo.jpeg"
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+    img = Image.open(img_name)
+    img_t = transform(img)
+    print(img_t.shape)  # pyright: ignore
+    batch_t = torch.unsqueeze(img_t, 0)  # pyright:ignore
+
+    # Predict
+    with torch.no_grad():
+        output = resnet_torch(batch_t)
+        print(output.shape)  # pyright: ignore
+        _, predicted = torch.max(output, 1)
+
+    print(
+        f"Predicted: {predicted.item()}"
+    )  # 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:
+        imagenet_labels = json.loads(url.read().decode())
+
+    label = imagenet_labels[str(predicted.item())][1]
+    print(f"Label for index {predicted.item()}: {label}")
+
+    identity = lambda x: x
+    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))
+
+    jax_batch = jnp.array(batch_t.numpy())
+    out, state = eqx.filter_vmap(
+        model, in_axes=(0, None), out_axes=(0, None), axis_name="batch"
+    )(jax_batch, state)
+    print(f"{out.shape}")
+
+    label = imagenet_labels[str(jnp.argmax(out))][1]
+    print(f"Label for index {jnp.argmax(out)}: {label}")
+
+
+if __name__ == "__main__":
+    test_resnet()