Skip to main content

Framework-agnostic library for checking array shapes at runtime.

Project description

JAX ResNet - Implementations and Checkpoints for ResNet Variants

Build & Tests

A Flax (Linen) implementation of ResNet (He et al. 2015), Wide ResNet (Zagoruyko & Komodakis 2016), ResNeXt (Xie et al. 2017), ResNet-D (He et al. 2020), and ResNeSt (Zhang et al. 2020). The code is modular so you can mix and match the various stem, residual, and bottleneck implementations.


You can install this package from PyPI:

pip install jax-resnet

Or directly from GitHub:

pip install --upgrade git+


See the bottom of jax-resnet/ for the available aliases/options for the ResNet variants (all models are in Flax)

Pretrained checkpoints from torch.hub are available for the following networks:

  • ResNet [18, 34, 50, 101, 152]
  • WideResNet [50, 101]
  • ResNeXt [50, 101]
  • ResNeSt [50-Fast, 50, 101, 200, 269]

The models are tested to have the same intermediate activations and outputs as the torch.hub implementations, except ResNeSt-50 Fast, whose activations don't match exactly but the final accuracy does.

A pretrained checkpoint for ResNetD-50 is available from The activations do not match exactly, but the final accuracy matches.

import jax.numpy as jnp
from jax_resnet import pretrained_resnest

ResNeSt50, variables = pretrained_resnest(50)
model = ResNeSt50()
out = model.apply(variables,
                  jnp.ones((32, 224, 224, 3)),  # ImageNet sized inputs.
                  mutable=False)  # Ensure `batch_stats` aren't updated.

You must install PyTorch yourself (instructions) to use these functions.

Transfer Learning

To extract a subset of the model, you can use Sequential(model.layers[start:end]).

The slice_variables function (found in in allows you to extract the corresponding subset of the variables dict. Check out that docstring for more information.

Checkpoint Accuracies

The top 1 and top 5 accuracies reported below are on the ImageNet2012 validation split. The data was preprocessed as in the official PyTorch example.

Model Size Top 1 Top 5
ResNet 18 69.75% 89.06%
34 73.29% 91.42%
50 76.13% 92.86%
101 77.37% 93.53%
152 78.30% 94.04%
Wide ResNet 50 78.48% 94.08%
101 78.88% 94.29%
ResNeXt 50 77.60% 93.70%
101 79.30% 94.51%
ResNet-D 50 77.57% 93.85%

The ResNeSt validation data was preprocessed as in zhang1989/ResNeSt.

Model Size Crop Size Top 1 Top 5
ResNeSt-Fast 50 224 80.53% 95.34%
ResNeSt 50 224 81.05% 95.42%
101 256 82.82% 96.32%
200 320 83.84% 96.86%
269 416 84.53% 96.98%


Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for jax-resnet, version 0.0.4
Filename, size File type Python version Upload date Hashes
Filename, size jax_resnet-0.0.4-py2.py3-none-any.whl (11.9 kB) File type Wheel Python version py2.py3 Upload date Hashes View

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring DigiCert DigiCert EV certificate Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page