update nfnets

leondgarse · Sep 24, 2021 · a9a05bf · a9a05bf
1 parent d51e218
commit a9a05bf
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -22,6 +22,7 @@
   - [ResMLP](#resmlp)
   - [GMLP](#gmlp)
   - [LeViT](#levit)
+  - [NFNets](#nfnets)
 - [Other implemented keras models](#other-implemented-keras-models)
 
 <!-- /TOC -->
@@ -253,11 +254,23 @@
   | LeViT192  | 11M    | 224              | 80.0     | [levit192_imagenet.h5](https://github.com/leondgarse/keras_cv_attention_models/releases/download/levit/levit192_imagenet.h5) |
   | LeViT256  | 19M    | 224              | 81.6     | [levit256_imagenet.h5](https://github.com/leondgarse/keras_cv_attention_models/releases/download/levit/levit256_imagenet.h5) |
   | LeViT384  | 39M    | 224              | 82.6     | [levit384_imagenet.h5](https://github.com/leondgarse/keras_cv_attention_models/releases/download/levit/levit384_imagenet.h5) |
+## NFNets
+  - [Keras NFNets](https://github.com/leondgarse/keras_cv_attention_models/tree/main/keras_cv_attention_models/nfnets) is for [PDF 2102.06171 High-Performance Large-Scale Image Recognition Without Normalization](https://arxiv.org/pdf/2102.06171.pdf).
+
+  | Model       | Params | Image  resolution | Top1 Acc | Download |
+  | ----------- | ------ | ----------------- | -------- | -------- |
+  | NFNetF0     | 71.5M  | 256               | 83.6     |          |
+  | NFNetF1     | 132.6M | 320               | 84.7     |          |
+  | NFNetF2     | 193.8M | 352               | 85.1     |          |
+  | NFNetF3     | 254.9M | 416               | 85.7     |          |
+  | NFNetF4     | 316.1M | 512               | 85.9     |          |
+  | NFNetF5     | 377.2M | 544               | 86.0     |          |
+  | NFNetF6 SAM | 438.4M | 576               | 86.5     |          |
+  | NFNetF7     | 499.5M |                   |          |          |
 ***
 
 # Other implemented keras models
   - [perceiver_image_classification](https://keras.io/examples/vision/perceiver_image_classification/)
-  - [Github ypeleg/nfnets-keras](https://github.com/ypeleg/nfnets-keras)
   - [Github faustomorales/vit-keras](https://github.com/faustomorales/vit-keras)
   - [Github tensorflow/resnet_rs](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs)
   - [Github google-research/big_transfer](https://github.com/google-research/big_transfer)

diff --git a/keras_cv_attention_models/__init__.py b/keras_cv_attention_models/__init__.py
@@ -18,5 +18,6 @@
 from keras_cv_attention_models.resnet_family import resnet_quad
 from keras_cv_attention_models.resnet_family import resnet_deep
 from keras_cv_attention_models import volo
+from keras_cv_attention_models import nfnets
 from keras_cv_attention_models import download_and_load
 from keras_cv_attention_models import visualizing
diff --git a/keras_cv_attention_models/common_layers.py b/keras_cv_attention_models/common_layers.py
@@ -108,8 +108,8 @@ def se_module(inputs, se_ratio=0.25, activation="relu", use_bias=True, name=""):
     se = keras.layers.Conv2D(reduction, kernel_size=1, use_bias=use_bias, kernel_initializer=CONV_KERNEL_INITIALIZER, name=name and name + "1_conv")(se)
     se = activation_by_name(se, activation=activation, name=name)
     se = keras.layers.Conv2D(filters, kernel_size=1, use_bias=use_bias, kernel_initializer=CONV_KERNEL_INITIALIZER, name=name and name + "2_conv")(se)
-    se = keras.layers.Activation("sigmoid")(se)
-    return keras.layers.Multiply()([inputs, se])
+    se = activation_by_name(se, activation="sigmoid", name=name)
+    return keras.layers.Multiply(name=name and name + "out")([inputs, se])
 
 
 def drop_block(inputs, drop_rate=0, name=None):

diff --git a/keras_cv_attention_models/nfnets/nfnets.py b/keras_cv_attention_models/nfnets/nfnets.py
@@ -115,7 +115,7 @@ def block(
     deep = activation_by_name_with_gamma(deep, activation, name=name + "deep_3_")
     deep = std_conv2d_with_init(deep, expanded_filter, 1, strides=1, padding="VALID", name=name + "deep_4_")
     if se_ratio > 0:
-        deep = se_module(deep, se_ratio=se_ratio, activation=activation, use_bias=True, name=name + "se_")
+        deep = se_module(deep, se_ratio=se_ratio, activation="relu", use_bias=True, name=name + "se_")
         deep *= attn_gain
 
     deep = drop_block(deep, drop_rate)
@@ -162,6 +162,7 @@ def NormFreeNet(
     drop_connect_rate=0,
     classifier_activation="softmax",
     drop_rate=0,
+    pretrained="imagenet",
     model_name="nfnet",
     kwargs=None
 ):
@@ -204,26 +205,26 @@ def NormFreeNet(
     model = keras.models.Model(inputs, nn, name=model_name)
     return model
 
-def NFNetF0(input_shape=(256, 256, 3), num_classes=1000, activation="gelu", classifier_activation="softmax", drop_rate=0.2, **kwargs):
+def NFNetF0(input_shape=(256, 256, 3), num_classes=1000, activation="gelu", drop_rate=0.2, pretrained="imagenet", **kwargs):
     return NormFreeNet(num_blocks=[1, 2, 6, 3], model_name="nfnetf0", **locals(), **kwargs)
 
-def NFNetF1(input_shape=(320, 320, 3), num_classes=1000, activation="gelu", classifier_activation="softmax", drop_rate=0.3, **kwargs):
+def NFNetF1(input_shape=(320, 320, 3), num_classes=1000, activation="gelu", drop_rate=0.3, pretrained="imagenet", **kwargs):
     return NormFreeNet(num_blocks=[2, 4, 12, 6], model_name="nfnetf1", **locals(), **kwargs)
 
-def NFNetF2(input_shape=(352, 352, 3), num_classes=1000, activation="gelu", classifier_activation="softmax", drop_rate=0.4, **kwargs):
+def NFNetF2(input_shape=(352, 352, 3), num_classes=1000, activation="gelu", drop_rate=0.4, pretrained="imagenet", **kwargs):
     return NormFreeNet(num_blocks=[3, 6, 18, 9], model_name="nfnetf2", **locals(), **kwargs)
 
-def NFNetF3(input_shape=(416, 416, 3), num_classes=1000, activation="gelu", classifier_activation="softmax", drop_rate=0.4, **kwargs):
+def NFNetF3(input_shape=(416, 416, 3), num_classes=1000, activation="gelu", drop_rate=0.4, pretrained="imagenet", **kwargs):
     return NormFreeNet(num_blocks=[4, 8, 24, 12], model_name="nfnetf3", **locals(), **kwargs)
 
-def NFNetF4(input_shape=(512, 512, 3), num_classes=1000, activation="gelu", classifier_activation="softmax", drop_rate=0.5, **kwargs):
+def NFNetF4(input_shape=(512, 512, 3), num_classes=1000, activation="gelu", drop_rate=0.5, pretrained="imagenet", **kwargs):
     return NormFreeNet(num_blocks=[5, 10, 30, 15], model_name="nfnetf4", **locals(), **kwargs)
 
-def NFNetF5(input_shape=(544, 544, 3), num_classes=1000, activation="gelu", classifier_activation="softmax", drop_rate=0.5, **kwargs):
+def NFNetF5(input_shape=(544, 544, 3), num_classes=1000, activation="gelu", drop_rate=0.5, pretrained="imagenet", **kwargs):
     return NormFreeNet(num_blocks=[6, 12, 36, 18], model_name="nfnetf5", **locals(), **kwargs)
 
-def NFNetF6(input_shape=(576, 576, 3), num_classes=1000, activation="gelu", classifier_activation="softmax", drop_rate=0.5, **kwargs):
+def NFNetF6(input_shape=(576, 576, 3), num_classes=1000, activation="gelu", drop_rate=0.5, pretrained="imagenet", **kwargs):
     return NormFreeNet(num_blocks=[7, 14, 42, 21], model_name="nfnetf6", **locals(), **kwargs)
 
-def NFNetF7(input_shape=(608, 608, 3), num_classes=1000, activation="gelu", classifier_activation="softmax", drop_rate=0.5, **kwargs):
+def NFNetF7(input_shape=(608, 608, 3), num_classes=1000, activation="gelu", drop_rate=0.5, pretrained="imagenet", **kwargs):
     return NormFreeNet(num_blocks=[8, 16, 48, 24], model_name="nfnetf7", **locals(), **kwargs)