zibbini / convnext-v2_tensorflow

Tensorflow implementations of ConvNeXt V1 + V2 models w/ weights, including conversion and evaluation scripts.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

ConvNeXt V1 & V2

The ConvNeXt family of models are a series of convolutional neural networks (ConvNets) that achieve state-of-the-art results on image classification benchmarks, and can be readily applied to plenty of other image-based tasks. For more info on these models, refer to the following papers:

This repo contains unofficial implementations of both versions of ConvNeXt models in tensorflow. These are pretty much straight copies from the official PyTorch implementations (ConvNeXt V1 & ConvNeXt V2 with minor changes to suit Keras' functional API. Of note, these implementations:

  • Can accept custom input shapes and input tensors
  • Are capable of mixed precision use
  • Include model configurations (atto, femto, pico, nano) for ConvNeXt-V1 matching those found in the official pytorch-image-models library.

Weights are currently available for ConvNeXt V1 official configurations, but ConvNeXt V2 weights are still yet to be processed.

Usage

For installation of tensorflow and keras, please refer to the following guide. Once TensorFlow is installed, you can make use of the models like so:

# Construct 'tiny' model with standard ImageNet resolution
from convnext_tf import convnext_v1

model = convnext_v1.convnext_tiny(input_shape=(224, 224, 3))
# Construct the model with pre-processing layers
from tensorflow import keras
from convnext_tf import convnext_v1

inputs = keras.Input((224, 224, 3))
input_tensor = keras.layers.Rescaling(1 / 255)(inputs) # Rescale inputs to 0 - 1
cxn_tiny = convnext_v1.convnext_tiny(input_tensor=input_tensor)
model = keras.Model(inputs, cxn_tiny.output)
# Construct 'tiny' model with custom resolution
from convnext_tf import convnext_v1

model = convnext_v1.convnext_tiny(input_shape=(512, 512, 3))
# Use imagenet weights to load model
from convnext_tf import convnext_v1

# Load weights trained on ImageNet-1k dataset
model = convnext_v1.convnext_tiny(weights='imagenet_1k') 
# Load weights trained on ImageNet-22k dataset and fine-tuned on ImageNet-1k dataset
model = convnext_v1.convnext_tiny(weights='imagenet_22k') 
# Use custom classification head
from convnext_tf import convnext_v1
from tensorflow import keras

inputs = keras.Input((224, 224, 3))
cxn_tiny = convnext_v1.convnext_tiny(input_tensor=inputs, include_top=False)
x = cxn_tiny.output
x = keras.layers.GlobalAveragePooling2D()(x)
x = keras.layers.Dense(128, activation='relu')(x)
x = keras.layers.Dense(32, activation='relu')(x)
outputs = keras.layers.Dense(4, activation='sigmoid')(x)
model = keras.Model(inputs, outputs)

Conversion and Evaluation

Conversion scripts are convert_weights_v1.py and convert_weights_v2.py for ConvNeXt V1 and V2 models respectively. Evaluation code can be found in eval_1k.py. Note that the official validation ground truth labels found in the dev kit on ImageNet's website are incorrect - instead, you need to download the bounding box validation dataset from their website. The XML files in this dataset will include the correct ground truth labels which can be used for evaluation purposes.

Results and pre-trained models

These are comparison results between PyTorch and TensorFlow implementations from evaluation on the ImageNet validation dataset. Converted weights for use in TensorFlow are linked in the model columns.

ImageNet-1K trained models
name resolution PyTorch acc@1 TensorFlow acc@1 #params FLOPs model
ConvNeXt-T 224x224 82.1 81.3 28M 4.5G model
ConvNeXt-S 224x224 83.1 82.4 50M 8.7G model
ConvNeXt-B 224x224 83.8 83.3 89M 15.4G model
ConvNeXt-B 384x384 85.1 84.9 89M 45.0G model
ConvNeXt-L 224x224 84.3 83.9 198M 34.4G model
ConvNeXt-L 384x384 85.5 85.4 198M 101.0G model
ImageNet-22K trained models
name resolution PyTorch acc@1 TensorFlow acc@1 #params FLOPs model
ConvNeXt-T 224x224 82.9 82.3 29M 4.5G model
ConvNeXt-T 384x384 84.1 84.0 29M 13.1G model
ConvNeXt-S 224x224 84.6 84.1 50M 8.7G model
ConvNeXt-S 384x384 85.8 85.8 50M 25.5G model
ConvNeXt-B 224x224 85.8 85.4 89M 15.4G model
ConvNeXt-B 384x384 86.8 86.8 89M 47.0G model
ConvNeXt-L 224x224 86.6 86.4 198M 34.4G model
ConvNeXt-L 384x384 87.5 87.5 198M 101.0G model
ConvNeXt-XL 224x224 87.0 86.8 350M 60.9G model
ConvNeXt-XL 384x384 87.8 87.7 350M 179.0G model

Acknowledgements and References

About

Tensorflow implementations of ConvNeXt V1 + V2 models w/ weights, including conversion and evaluation scripts.

License:MIT License


Languages

Language:Python 100.0%