Equimo 0.1.0__tar.gz

equimo-0.1.0/Equimo.egg-info/PKG-INFO

@@ -0,0 +1,104 @@
+Metadata-Version: 2.1
+Name: Equimo
+Version: 0.1.0
+Summary: Implementation of popular vision models in Jax
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: einops>=0.8.0
+Requires-Dist: equinox>=0.11.5
+Requires-Dist: jax>=0.4.25
+Requires-Dist: jaxlib>=0.4.25
+
+# Equimo: Modern Vision Models in JAX/Equinox
+
+**WARNING**: This is a research library implementing recent computer vision models. The implementations are based on paper descriptions and may not be exact replicas of the original implementations. Use with caution in production environments.
+
+Equimo (Equinox Image Models) provides JAX/Equinox implementations of recent computer vision models, currently focusing (but not limited to) on transformer and state-space architectures.
+
+## Features
+
+- Pure JAX/Equinox implementations
+- Focus on recent architectures (2023-2024 papers)
+- Modular design for easy experimentation
+- Extensive documentation and type hints
+
+## Installation
+
+### From PyPI
+
+```bash
+pip install equimo
+```
+
+### From Source
+
+```bash
+git clone https://github.com/clementpoiret/equimo.git
+cd equimo
+pip install -e .
+```
+
+## Implemented Models
+
+| Model | Paper | Year | Status |
+|-------|-------|------|--------|
+| FasterViT | [FasterViT: Fast Vision Transformers with Hierarchical Attention](https://arxiv.org/abs/2306.06189) | 2023 | ✅ |
+| Castling-ViT | [Castling-ViT: Compressing Self-Attention via Switching Towards Linear-Angular Attention During Vision Transformer Inference](https://arxiv.org/abs/2211.10526) | 2023 | Partial* |
+| MLLA | [Mamba-like Linear Attention](https://arxiv.org/abs/2405.16605) | 2024 | ✅ |
+| PartialFormer | [Efficient Vision Transformers with Partial Attention](https://eccv.ecva.net/virtual/2024/poster/1877) | 2024 | ✅ |
+| SHViT | [SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design](https://arxiv.org/abs/2401.16456) | 2024 | ✅ |
+| VSSD | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559) | 2024 | ✅ |
+
+*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
+
+## Basic Usage
+
+```python
+import jax
+
+import equimo.models as em
+
+# Create a model (e.g. `faster_vit_0_224`)
+key = jax.random.PRNGKey(0)
+model = em.FasterViT(
+    img_size=224,
+    in_channels=3,
+    dim=64,
+    in_dim=64,
+    depths=[2, 3, 6, 5],
+    num_heads=[2, 4, 8, 16],
+    hat=[False, False, True, False],
+    window_size=[7, 7, 7, 7],
+    ct_size=2,
+    key=key,
+)
+
+# Generate random input
+x = jax.random.normal(key, (3, 224, 224))
+
+# Run inference
+output = model(x, enable_dropout=False, key=key)
+```
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## Citation
+
+If you use Equimo in your research, please cite:
+
+```bibtex
+@software{equimo2024,
+  author = {Clément POIRET},
+  title = {Equimo: Modern Vision Models in JAX/Equinox},
+  year = {2024},
+  publisher = {GitHub},
+  url = {https://github.com/clementpoiret/equimo}
+}
+```

equimo-0.1.0/Equimo.egg-info/SOURCES.txt

@@ -0,0 +1,32 @@
+LICENSE.md
+README.md
+pyproject.toml
+Equimo.egg-info/PKG-INFO
+Equimo.egg-info/SOURCES.txt
+Equimo.egg-info/dependency_links.txt
+Equimo.egg-info/requires.txt
+Equimo.egg-info/top_level.txt
+equimo/__init__.py
+equimo/utils.py
+equimo/layers/__init__.py
+equimo/layers/attention.py
+equimo/layers/convolution.py
+equimo/layers/downsample.py
+equimo/layers/dropout.py
+equimo/layers/ffn.py
+equimo/layers/generic.py
+equimo/layers/mamba.py
+equimo/layers/norm.py
+equimo/layers/patch.py
+equimo/layers/posemb.py
+equimo/layers/sharing.py
+equimo/layers/squeeze_excite.py
+equimo/models/__init__.py
+equimo/models/emamodel.py
+equimo/models/fastervit.py
+equimo/models/mlla.py
+equimo/models/partialformer.py
+equimo/models/shvit.py
+equimo/models/vit.py
+equimo/models/vssd.py
+equimo/ops/scan.py

equimo-0.1.0/Equimo.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

equimo-0.1.0/Equimo.egg-info/requires.txt

@@ -0,0 +1,4 @@
+einops>=0.8.0
+equinox>=0.11.5
+jax>=0.4.25
+jaxlib>=0.4.25

equimo-0.1.0/Equimo.egg-info/top_level.txt

@@ -0,0 +1 @@
+equimo

equimo-0.1.0/LICENSE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Clément POIRET
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

equimo-0.1.0/PKG-INFO

@@ -0,0 +1,104 @@
+Metadata-Version: 2.1
+Name: Equimo
+Version: 0.1.0
+Summary: Implementation of popular vision models in Jax
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: einops>=0.8.0
+Requires-Dist: equinox>=0.11.5
+Requires-Dist: jax>=0.4.25
+Requires-Dist: jaxlib>=0.4.25
+
+# Equimo: Modern Vision Models in JAX/Equinox
+
+**WARNING**: This is a research library implementing recent computer vision models. The implementations are based on paper descriptions and may not be exact replicas of the original implementations. Use with caution in production environments.
+
+Equimo (Equinox Image Models) provides JAX/Equinox implementations of recent computer vision models, currently focusing (but not limited to) on transformer and state-space architectures.
+
+## Features
+
+- Pure JAX/Equinox implementations
+- Focus on recent architectures (2023-2024 papers)
+- Modular design for easy experimentation
+- Extensive documentation and type hints
+
+## Installation
+
+### From PyPI
+
+```bash
+pip install equimo
+```
+
+### From Source
+
+```bash
+git clone https://github.com/clementpoiret/equimo.git
+cd equimo
+pip install -e .
+```
+
+## Implemented Models
+
+| Model | Paper | Year | Status |
+|-------|-------|------|--------|
+| FasterViT | [FasterViT: Fast Vision Transformers with Hierarchical Attention](https://arxiv.org/abs/2306.06189) | 2023 | ✅ |
+| Castling-ViT | [Castling-ViT: Compressing Self-Attention via Switching Towards Linear-Angular Attention During Vision Transformer Inference](https://arxiv.org/abs/2211.10526) | 2023 | Partial* |
+| MLLA | [Mamba-like Linear Attention](https://arxiv.org/abs/2405.16605) | 2024 | ✅ |
+| PartialFormer | [Efficient Vision Transformers with Partial Attention](https://eccv.ecva.net/virtual/2024/poster/1877) | 2024 | ✅ |
+| SHViT | [SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design](https://arxiv.org/abs/2401.16456) | 2024 | ✅ |
+| VSSD | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559) | 2024 | ✅ |
+
+*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
+
+## Basic Usage
+
+```python
+import jax
+
+import equimo.models as em
+
+# Create a model (e.g. `faster_vit_0_224`)
+key = jax.random.PRNGKey(0)
+model = em.FasterViT(
+    img_size=224,
+    in_channels=3,
+    dim=64,
+    in_dim=64,
+    depths=[2, 3, 6, 5],
+    num_heads=[2, 4, 8, 16],
+    hat=[False, False, True, False],
+    window_size=[7, 7, 7, 7],
+    ct_size=2,
+    key=key,
+)
+
+# Generate random input
+x = jax.random.normal(key, (3, 224, 224))
+
+# Run inference
+output = model(x, enable_dropout=False, key=key)
+```
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## Citation
+
+If you use Equimo in your research, please cite:
+
+```bibtex
+@software{equimo2024,
+  author = {Clément POIRET},
+  title = {Equimo: Modern Vision Models in JAX/Equinox},
+  year = {2024},
+  publisher = {GitHub},
+  url = {https://github.com/clementpoiret/equimo}
+}
+```

equimo-0.1.0/README.md

@@ -0,0 +1,92 @@
+# Equimo: Modern Vision Models in JAX/Equinox
+
+**WARNING**: This is a research library implementing recent computer vision models. The implementations are based on paper descriptions and may not be exact replicas of the original implementations. Use with caution in production environments.
+
+Equimo (Equinox Image Models) provides JAX/Equinox implementations of recent computer vision models, currently focusing (but not limited to) on transformer and state-space architectures.
+
+## Features
+
+- Pure JAX/Equinox implementations
+- Focus on recent architectures (2023-2024 papers)
+- Modular design for easy experimentation
+- Extensive documentation and type hints
+
+## Installation
+
+### From PyPI
+
+```bash
+pip install equimo
+```
+
+### From Source
+
+```bash
+git clone https://github.com/clementpoiret/equimo.git
+cd equimo
+pip install -e .
+```
+
+## Implemented Models
+
+| Model | Paper | Year | Status |
+|-------|-------|------|--------|
+| FasterViT | [FasterViT: Fast Vision Transformers with Hierarchical Attention](https://arxiv.org/abs/2306.06189) | 2023 | ✅ |
+| Castling-ViT | [Castling-ViT: Compressing Self-Attention via Switching Towards Linear-Angular Attention During Vision Transformer Inference](https://arxiv.org/abs/2211.10526) | 2023 | Partial* |
+| MLLA | [Mamba-like Linear Attention](https://arxiv.org/abs/2405.16605) | 2024 | ✅ |
+| PartialFormer | [Efficient Vision Transformers with Partial Attention](https://eccv.ecva.net/virtual/2024/poster/1877) | 2024 | ✅ |
+| SHViT | [SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design](https://arxiv.org/abs/2401.16456) | 2024 | ✅ |
+| VSSD | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559) | 2024 | ✅ |
+
+*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
+
+## Basic Usage
+
+```python
+import jax
+
+import equimo.models as em
+
+# Create a model (e.g. `faster_vit_0_224`)
+key = jax.random.PRNGKey(0)
+model = em.FasterViT(
+    img_size=224,
+    in_channels=3,
+    dim=64,
+    in_dim=64,
+    depths=[2, 3, 6, 5],
+    num_heads=[2, 4, 8, 16],
+    hat=[False, False, True, False],
+    window_size=[7, 7, 7, 7],
+    ct_size=2,
+    key=key,
+)
+
+# Generate random input
+x = jax.random.normal(key, (3, 224, 224))
+
+# Run inference
+output = model(x, enable_dropout=False, key=key)
+```
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## Citation
+
+If you use Equimo in your research, please cite:
+
+```bibtex
+@software{equimo2024,
+  author = {Clément POIRET},
+  title = {Equimo: Modern Vision Models in JAX/Equinox},
+  year = {2024},
+  publisher = {GitHub},
+  url = {https://github.com/clementpoiret/equimo}
+}
+```