Root package info.
Project description
Eqxvision
Eqxvision is a package of popular computer vision model architectures built using Equinox.
Installation
Use the package manager pip to install eqxvision.
pip install eqxvision
requires: python>=3.7
Usage
Picking a model and doing a forward pass is as simple as ...
import jax
import jax.random as jr
import equinox as eqx
from eqxvision.models import alexnet
@eqx.filter_jit
def forward(net, images, key):
keys = jax.random.split(key, images.shape[0])
output = jax.vmap(net, axis_name=('batch'))(images, key=keys)
...
net = alexnet(num_classes=1000)
images = jr.uniform(jr.PRNGKey(0), shape=(1,3,224,224))
output = forward(net, images, jr.PRNGKey(0))
What's New?
-
[Experimental]
Now supports loading PyTorch weights fromtorchvision
for models without BatchNorm!!! note Due to slight differences in the implementation of underlying operations, differences in the output values can be expected from
torchvision
models.
Tips
- Better to use
@equinox.jit_filter
instead of@jax.jit
- Advisable to use
jax.{v,p}map
withaxis_name='batch'
for all models - Don't forget to switch to
inference
mode for evaluations - Wrap with
eqx.filter(net, eqx.is_array)
forOptax
initialisation.
Contributing
Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.
Development Process
If you plan to modify the code or documentation, please follow the steps below:
- Fork the repository and create your branch from
dev
. - If you have modified the code (new feature or bug-fix), please add unit tests.
- If you have changed APIs, update the documentation. Make sure the documentation builds.
mkdocs serve
- Ensure the test suite passes.
pytest tests -vvv
- Make sure your code passes the formatting checks. Automatically checked with a
pre-commit
hook.
Acknowledgements
License
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for eqxvision-0.1.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 83e9820aea7f19a19b6144dc5ae1a990c2b8a99075f23d2b2ec8888ce16557e2 |
|
MD5 | 2394de0106dae4b87206a5156c217952 |
|
BLAKE2b-256 | c81b3698bb5f590810d1d689a47e955dd33d12c275a4955bac68887a107b7245 |