eqxvision 0.0.4a0

Root package info.

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

Usage

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)(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))

import equinox as eqx
from eqxvision.models import alexnet

net = alexnet(num_classes=1000)
net = eqx.tree_inference(net, True)

Tips

• Checkout the documentation for a sample usage for each model;
• Better to use @equinox.jit_filter instead of @jax.jit;
• Use jax.vmap(net, axis_name='batch')(images) for models with batchnorms;
• Don't forget to switch to inference mode for evaluations.

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