https://github.com/rolandgvc/flaxvision

A selection of neural network models ported from torchvision for JAX & Flax.
https://github.com/rolandgvc/flaxvision

Last synced: 3 months ago
JSON representation

A selection of neural network models ported from torchvision for JAX & Flax.

• Host: GitHub
• URL: https://github.com/rolandgvc/flaxvision
• Owner: rolandgvc
• License: apache-2.0
• Created: 2020-06-14T16:34:14.000Z (over 4 years ago)
• Default Branch: master
• Last Pushed: 2021-01-04T13:49:58.000Z (almost 4 years ago)
• Last Synced: 2024-04-11T00:42:47.189Z (7 months ago)
• Language: Python
• Homepage:
• Size: 168 KB
• Stars: 43
• Watchers: 4
• Forks: 2
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • Changelog: CHANGELOG.md
  • Contributing: CONTRIBUTING.md
  • License: LICENSE

Awesome Lists containing this project

• awesome-jax - FlaxVision - Flax version of TorchVision. <img src="https://img.shields.io/github/stars/rolandgvc/flaxvision?style=social" align="center"> (Libraries / New Libraries)

README

        
# flaxvision

The flaxvision package contains a selection of neural network models ported from [torchvision](https://github.com/pytorch/vision) to be used with [JAX](https://github.com/google/jax) & [Flax](https://github.com/google/flax).

**Note: flaxvision is currently in active development. API and functionality may change between releases.**

### Roadmap to v0.1.0

Planned features for the first release:

- [x] Update models to [linen API](https://github.com/google/flax/tree/0132b3f234a9868b47df491efde870bdc58e97a9/linen_examples)

- [x] Add support for transfer learning

- [x] Add support to ResNet for dilated convolutions

- [x] Port DeepLabv3 model for image segmentation

## Quickstart

### Transfer Learning Example

```python

from jax import random

from flaxvision import models

rng = random.PRNGKey(0)

pretrained_model = models.vgg16(rng, pretrained=True)

```

## How To Contribute

If interested in adding additional models or improving existent ones, please start by openning an Issue describing your idea.

## Acknowledgments

The initial work for flaxvision started during the Google Summer of Code program at Google AI under [Avital Oliver](https://github.com/avital)'s mentorship.