A few more simple updates.

Author: mingxingtan

efficientdet/README.md

@@ -10,6 +10,7 @@ Arxiv link: https://arxiv.org/abs/1911.09070
 
 Updates:
 
+  - Jul19/2021: Added Nvidia TensorRT script/instruction [link](https://github.com/NVIDIA/TensorRT/tree/master/samples/python/efficientdet)
   - May10/2021: Added EfficientDet-lite checkpoints (by Yuqi and TFLite team)
   - Mar25/2021: Added [Det-AdvProp](https://arxiv.org/abs/2103.13886) model checkpoints ([see this page](./Det-AdvProp.md)).
   - Jul20/2020: Added keras/TF2 and new SOTA D7x: 55.1mAP with 153ms.
@@ -77,6 +78,26 @@ We have provided a list of EfficientDet checkpoints and results as follows:
 
 For more accurate and robust EfficientDet, please see [this page](./Det-AdvProp.md), which contains a list of models trained with Det-AdvProp + AutoAugment (AA) described in [this paper](https://arxiv.org/abs/2103.13886). The obatined model is not only more accurate on clean images, but also much more robust against various corruptions and domain shift.
 
+On single Tesla V100 without using TensorRT, our end-to-end
+latency and throughput are:
+
+
+|       Model    |   mAP | batch1 latency |  batch1 throughput |  batch8 throughput |
+| ------ | ------ | ------  | ------ | ------ |
+| EfficientDet-D0 |  34.6 | 10.2ms | 97 fps | 209 fps |
+| EfficientDet-D1 |  40.5 | 13.5ms | 74 fps | 140 fps |
+| EfficientDet-D2 |  43.0 | 17.7ms | 57 fps | 97 fps  |
+| EfficientDet-D3 |  47.5 | 28.0ms | 36 fps | 58 fps  |
+| EfficientDet-D4 |  49.7 | 42.8ms | 23 fps | 35 fps  |
+| EfficientDet-D5 |  51.5 | 72.5ms | 14 fps | 18 fps  |
+| EfficientDet-D6 |  52.6 | 92.8ms | 11 fps | - fps  |
+| EfficientDet-D7 |  53.7 | 122ms  | 8.2 fps | - fps  |
+| EfficientDet-D7x |  55.1 | 153ms  | 6.5 fps | - fps  |
+
+** FPS means frames per second (or images/second).
+
+** EfficientDet can be significantly sped up with TensorRT: [link](https://github.com/NVIDIA/TensorRT/tree/master/samples/python/efficientdet)
+
 In addition, the following table includes a list of models trained with fixed 640x640 image sizes (see appendix of [this paper](https://arxiv.org/abs/1911.09070)):
 
 
@@ -152,24 +173,6 @@ use the following command:
       --model_name=efficientdet-d0  --input_image=testdata/img1.jpg  \
       --output_image_dir=/tmp/
 
-On single Tesla V100 without using TensorRT, our end-to-end
-latency and throughput are:
-
-
-|       Model    |   mAP | batch1 latency |  batch1 throughput |  batch8 throughput |
-| ------ | ------ | ------  | ------ | ------ |
-| EfficientDet-D0 |  34.6 | 10.2ms | 97 fps | 209 fps |
-| EfficientDet-D1 |  40.5 | 13.5ms | 74 fps | 140 fps |
-| EfficientDet-D2 |  43.0 | 17.7ms | 57 fps | 97 fps  |
-| EfficientDet-D3 |  47.5 | 28.0ms | 36 fps | 58 fps  |
-| EfficientDet-D4 |  49.7 | 42.8ms | 23 fps | 35 fps  |
-| EfficientDet-D5 |  51.5 | 72.5ms | 14 fps | 18 fps  |
-| EfficientDet-D6 |  52.6 | 92.8ms | 11 fps | - fps  |
-| EfficientDet-D7 |  53.7 | 122ms  | 8.2 fps | - fps  |
-| EfficientDet-D7x |  55.1 | 153ms  | 6.5 fps | - fps  |
-
-** FPS means frames per second (or images/second).
-
 ## 5. Inference for images.
 
     # Step0: download model and testing image.

efficientnetv2/README.md

@@ -5,6 +5,11 @@
 [TF-hub![TF-Hub In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.sandbox.google.com/github/google/automl/blob/master/efficientnetv2/tfhub.ipynb)
 
 
+   - Jul19/2021: A list of updates
+     * Added TF2 script [here](https://github.com/google/automl/blob/master/efficientnetv2/main_tf2.py).
+     * Updated ImageNet21k sigmoid-loss checkpoints, for multi-class pseudo labeling.
+     * Add EfficientNetV2-XL 21k and 1k checkpoint and hub modules.
+     * Added Nvidia TensorRT script [here](https://github.com/NVIDIA/TensorRT/tree/master/samples/python/efficientnet).
    - May13/2021: Initial code release for [EfficientNetV2 models](https://arxiv.org/abs/2104.00298): accepted to ICML'21.
 
 ## 1. About EfficientNetV2 Models
@@ -28,7 +33,7 @@ We have provided a list of results and checkpoints as follows:
 |    EffNetV2-M     |    85.2%   |    54.1M    | 24.7B    | [V100/A100](g3doc/effnetv2-m-gpu.png) |  [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-m.tgz),  [tensorboard](https://tensorboard.dev/experiment/syoaqB2gTP6Vr0KRlrezmg)
 |    EffNetV2-L     |    85.7%   |   119.5M    | 56.3B    | [V100/A100](g3doc/effnetv2-l-gpu.png) |  [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-l.tgz),  [tensorboard](https://tensorboard.dev/experiment/qgnTQ5JZQ92nSex6ZlWBbQ)
 
-** Thanks NVIDIA for providing the inference latency (benchmark scripts coming soon)
+** Thanks NVIDIA for providing the inference latency: full TensorRT scripts and instructions are available here: [link](https://github.com/NVIDIA/TensorRT/tree/master/samples/python/efficientnet)
 
 
 Here are a list of ImageNet21K pretrained and finetuned models:
@@ -39,6 +44,7 @@ Here are a list of ImageNet21K pretrained and finetuned models:
 |  EffNetV2-S   |  [pretrain ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-s-21k.tgz)  |  top1=84.9%,  [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-s-21k-ft1k.tgz),  [tensorboard](https://tensorboard.dev/experiment/7sga2olqTBeH4ioydel0hg/) |
 |  EffNetV2-M   |  [pretrain ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-m-21k.tgz)  |  top1=86.2%,  [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-m-21k-ft1k.tgz),  [tensorboard](https://tensorboard.dev/experiment/HkV6ANZSQ6WI5GhlZa48xQ/) |
 |  EffNetV2-L   |  [pretrain ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-l-21k.tgz)  |  top1=86.9%,  [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-l-21k-ft1k.tgz),  [tensorboard](https://tensorboard.dev/experiment/m9ZHx1L6SQu5iBYhXO5jOw/) |
+|  EffNetV2-XL   |  [pretrain ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-xl-21k.tgz)  |  top1=87.2%,  [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-xl-21k-ft1k.tgz),  [tensorboard]()|
 
 For comparison with EfficientNetV1, we have also provided a few smaller V2 models using the same scaling and preprocessing as V1:
 

efficientnetv2/effnetv2_model.py

@@ -25,11 +25,10 @@
 import copy
 import itertools
 import math
+import os
 
 from absl import logging
 import numpy as np
-import six
-from six.moves import xrange
 import tensorflow as tf
 
 import effnetv2_configs
@@ -560,7 +559,7 @@ def _build(self):
         block_args.input_filters = block_args.output_filters
         block_args.strides = 1
         # pylint: enable=protected-access
-      for _ in xrange(block_args.num_repeat - 1):
+      for _ in range(block_args.num_repeat - 1):
         self._blocks.append(
             conv_block(block_args, self._mconfig, name=block_name()))
 
@@ -624,7 +623,7 @@ def call(self, inputs, training, with_endpoints=False):
       if is_reduction:
         self.endpoints['reduction_%s' % reduction_idx] = outputs
       if block.endpoints:
-        for k, v in six.iteritems(block.endpoints):
+        for k, v in block.endpoints.items():
           self.endpoints['block_%s/%s' % (idx, k)] = v
           if is_reduction:
             self.endpoints['reduction_%s/%s' % (reduction_idx, k)] = v
@@ -655,7 +654,7 @@ def call(self, inputs, training, with_endpoints=False):
 def get_model(model_name,
               model_config=None,
               include_top=True,
-              pretrained=True,
+              weights='imagenet',
               training=True,
               with_endpoints=False,
               **kwargs):
@@ -667,7 +666,12 @@ def get_model(model_name,
     model_name: a string such as 'efficientnetv2-s' or 'efficientnet-b0'.
     model_config: A dict of model configurations or a string of hparams.
     include_top: whether to include the final dense layer for classification.
-    pretrained: if true, download the checkpoint. If string, load the ckpt.
+    weights: One of None (random initialization),
+      'imagenet' (pretrained on ImageNet),
+      'imagenet21k' (pretrained on Imagenet21k),
+      'imagenet21k-ft1k' (pretrained on 21k and finetuned on 1k), 
+      'jft' (trained with non-labelled JFT-300),
+      or the path to the weights file to be loaded. Defaults to 'imagenet'.
     training: If true, all model variables are trainable.
     with_endpoints: whether to return all intermedia endpoints.
     **kwargs: additional parameters for keras model, such as name=xx.
@@ -679,27 +683,106 @@ def get_model(model_name,
   net(tf.keras.Input(shape=(None, None, 3)),
       training=training,
       with_endpoints=with_endpoints)
-  if pretrained is True:  # pylint: disable=g-bool-id-comparison
-    # pylint: disable=line-too-long
-    # download checkpoint and set pretrained path. Supported models include:
-    #   efficientnetv2-s, efficientnetv2-m, efficientnetv2-l,
-    #   efficientnetv2-b0, efficientnetv2-b1, efficientnetv2-b2, efficientnetv2-b3,
-    #   efficientnet-b0, efficientnet-b1, efficientnet-b2, efficientnet-b3,
-    #   efficientnet-b4, efficientnet-b5, efficientnet-b6, efficientnet-b7,
-    #   efficientnet-l2
-    # v2: https://github.com/google/automl/tree/master/efficientnetv2
-    # v1: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet
-    # pylint: enable=line-too-long
-
-    url = ('https://storage.googleapis.com/cloud-tpu-checkpoints/'
-           f'efficientnet/v2/{model_name}.tgz')
-    pretrained_ckpt = tf.keras.utils.get_file(model_name, url, untar=True)
-  else:
-    pretrained_ckpt = pretrained
 
-  if pretrained_ckpt:
-    if tf.io.gfile.isdir(pretrained_ckpt):
-      pretrained_ckpt = tf.train.latest_checkpoint(pretrained_ckpt)
-    net.load_weights(pretrained_ckpt)
+  if not weights:  # pylint: disable=g-bool-id-comparison
+    return net
+
+  v2url = 'https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/'
+  v1url = 'https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/'
+  v1jfturl = 'https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/'
+  pretrained_ckpts = {
+      # EfficientNet V2.
+      'efficientnetv2-s': {
+          'imagenet': v2url + 'efficientnetv2-s.tgz',
+          'imagenet21k': v2url + 'efficientnetv2-s-21k.tgz',
+          'imagenet21k-ft1k': v2url + 'efficientnetv2-s-21k-ft1k.tgz',
+      },
+      'efficientnetv2-m': {
+          'imagenet': v2url + 'efficientnetv2-m.tgz',
+          'imagenet21k': v2url + 'efficientnetv2-m-21k.tgz',
+          'imagenet21k-ft1k': v2url + 'efficientnetv2-m-21k-ft1k.tgz',
+      },
+      'efficientnetv2-l': {
+          'imagenet': v2url + 'efficientnetv2-l.tgz',
+          'imagenet21k': v2url + 'efficientnetv2-l-21k.tgz',
+          'imagenet21k-ft1k': v2url + 'efficientnetv2-l-21k-ft1k.tgz',
+      },
+      'efficientnetv2-xl': {
+          # no imagenet ckpt.
+          'imagenet21k': v2url + 'efficientnetv2-xl-21k.tgz',
+          'imagenet21k-ft1k': v2url + 'efficientnetv2-xl-21k-ft1k.tgz',
+      },
+
+      'efficientnetv2-b0': {
+          'imagenet': v2url + 'efficientnetv2-b0.tgz',
+          'imagenet21k': v2url + 'efficientnetv2-b0-21k.tgz',
+          'imagenet21k-ft1k': v2url + 'efficientnetv2-b0-21k-ft1k.tgz',
+      },
+      'efficientnetv2-b1': {
+          'imagenet': v2url + 'efficientnetv2-b1.tgz',
+          'imagenet21k': v2url + 'efficientnetv2-b1-21k.tgz',
+          'imagenet21k-ft1k': v2url + 'efficientnetv2-b1-21k-ft1k.tgz',
+      },
+      'efficientnetv2-b2': {
+          'imagenet': v2url + 'efficientnetv2-b2.tgz',
+          'imagenet21k': v2url + 'efficientnetv2-b2-21k.tgz',
+          'imagenet21k-ft1k': v2url + 'efficientnetv2-b2-21k-ft1k.tgz',
+      },
+      'efficientnetv2-b3': {
+          'imagenet': v2url + 'efficientnetv2-b3.tgz',
+          'imagenet21k': v2url + 'efficientnetv2-b3-21k.tgz',
+          'imagenet21k-ft1k': v2url + 'efficientnetv2-b3-21k-ft1k.tgz',
+      },
+
+      # EfficientNet V1.
+      'efficientnet-b0': {
+          'imagenet': v1url + 'efficientnet-b0.tar.gz',
+          'jft': v1jfturl + 'noisy_student_efficientnet-b0.tar.gz',
+      },
+      'efficientnet-b1': {
+          'imagenet': v1url + 'efficientnet-b1.tar.gz',
+          'jft': v1jfturl + 'noisy_student_efficientnet-b1.tar.gz',
+      },
+      'efficientnet-b2': {
+          'imagenet': v1url + 'efficientnet-b2.tar.gz',
+          'jft': v1jfturl + 'noisy_student_efficientnet-b2.tar.gz',
+      },
+      'efficientnet-b3': {
+          'imagenet': v1url + 'efficientnet-b3.tar.gz',
+          'jft': v1jfturl + 'noisy_student_efficientnet-b3.tar.gz',
+      },
+      'efficientnet-b4': {
+          'imagenet': v1url + 'efficientnet-b4.tar.gz',
+          'jft': v1jfturl + 'noisy_student_efficientnet-b4.tar.gz',
+      },
+      'efficientnet-b5': {
+          'imagenet': v1url + 'efficientnet-b5.tar.gz',
+          'jft': v1jfturl + 'noisy_student_efficientnet-b5.tar.gz',
+      },
+      'efficientnet-b6': {
+          'imagenet': v1url + 'efficientnet-b6.tar.gz',
+          'jft': v1jfturl + 'noisy_student_efficientnet-b6.tar.gz',
+      },
+      'efficientnet-b7': {
+          'imagenet': v1url + 'efficientnet-b7.tar.gz',
+          'jft': v1jfturl + 'noisy_student_efficientnet-b7.tar.gz',
+      },
+      'efficientnet-b8': {
+          'imagenet': v1url + 'efficientnet-b8.tar.gz',
+      },
+      'efficientnet-l2': {
+          'jft': v1jfturl + 'noisy_student_efficientnet-l2_475.tar.gz',
+      },
+  }
+
+  if model_name in pretrained_ckpts and weights in pretrained_ckpts[model_name]:
+    url = pretrained_ckpts[model_name][weights]
+    fname = os.path.splitext(os.path.basename(url))[0]
+    pretrained_ckpt= tf.keras.utils.get_file(fname, url , untar=True)
+  else:
+    pretrained_ckpt = weights
 
+  if tf.io.gfile.isdir(pretrained_ckpt):
+    pretrained_ckpt = tf.train.latest_checkpoint(pretrained_ckpt)
+  net.load_weights(pretrained_ckpt)
   return net