eqxvision 0.0.1

Root package info.

Eqxvision

Eqxvision is a Python library providing computer vision models to the Equinox ecosystem.

Installation

Use the package manager pip to install foobar.

pip install eqxvision

Usage

import jax
import jax.random as jr
import equinox as eqx
from eqxvision.models import resnet18

@eqx.filter_jit
def forward(net, images, key):
    keys = jax.random.split(key, images.shape[0])
    jax.vmap(net)(images, key=keys)

net = resnet18(num_classes=1000)

images = jr.uniform(jr.PRNGKey(0), shape=(1,3,224,224))
output = forward(net, images, jr.PRNGKey(0))

Tips

• Use jax.vmap(net, axis_name='batch')(images) for models with batchnorms.
• Don't forget to call eqx.inference for switching to inference mode.

Roadmap

• Add VGGs, Inception, GoogLeNet
• Add/Explore functionality to load weights directly from torch.pth
• Doc fixes
• Build fixes
• Pre-commit Hooks

Contributing

Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.

Please make sure to update tests as appropriate.

Acknowledgements

• Torchvision
• Equinox
• Patrick Kidger

License

MIT