├── dependencies.txt ├── efficientnet_tf ├── g3doc │ ├── flops.png │ ├── params.png │ ├── condconv-layer.png │ ├── lite-float-gpu.png │ ├── lite-quant-cpu.png │ ├── lite-quant-size.png │ └── efficientnet-edgetpu.png ├── __pycache__ │ ├── utils.cpython-37.pyc │ ├── eval_ckpt_main.cpython-37.pyc │ ├── lars_optimizer.cpython-37.pyc │ ├── preprocessing.cpython-37.pyc │ ├── efficientnet_builder.cpython-37.pyc │ ├── efficientnet_model.cpython-37.pyc │ └── model_builder_factory.cpython-37.pyc ├── tpu │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ └── efficientnet_x_builder.cpython-37.pyc │ ├── __init__.py │ └── efficientnet_x_builder.py ├── lite │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ └── efficientnet_lite_builder.cpython-37.pyc │ ├── __init__.py │ ├── efficientnet_lite_builder_test.py │ ├── README.md │ └── efficientnet_lite_builder.py ├── condconv │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── condconv_layers.cpython-37.pyc │ │ └── efficientnet_condconv_builder.cpython-37.pyc │ ├── __init__.py │ ├── README.md │ ├── efficientnet_condconv_builder.py │ └── condconv_layers.py ├── edgetpu │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ └── efficientnet_edgetpu_builder.cpython-37.pyc │ ├── __init__.py │ ├── README.md │ └── efficientnet_edgetpu_builder.py ├── __init__.py ├── model_builder_factory.py ├── inspect_model_architecture.py ├── eval_ckpt_main.py ├── export_model.py ├── lars_optimizer.py ├── preprocessing.py ├── README.md ├── efficientnet_builder.py ├── imagenet_input.py └── utils.py ├── temp.sh ├── requirements.txt ├── .gitignore ├── readme.md ├── get_efn_weights.py └── load_efficientnet.py /dependencies.txt: -------------------------------------------------------------------------------- 1 | efficientnet 2 | tensorflow 3 | numpy 4 | tqdm -------------------------------------------------------------------------------- /efficientnet_tf/g3doc/flops.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/g3doc/flops.png -------------------------------------------------------------------------------- /efficientnet_tf/g3doc/params.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/g3doc/params.png -------------------------------------------------------------------------------- /efficientnet_tf/g3doc/condconv-layer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/g3doc/condconv-layer.png -------------------------------------------------------------------------------- /efficientnet_tf/g3doc/lite-float-gpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/g3doc/lite-float-gpu.png -------------------------------------------------------------------------------- /efficientnet_tf/g3doc/lite-quant-cpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/g3doc/lite-quant-cpu.png -------------------------------------------------------------------------------- /efficientnet_tf/g3doc/lite-quant-size.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/g3doc/lite-quant-size.png -------------------------------------------------------------------------------- /efficientnet_tf/__pycache__/utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/__pycache__/utils.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/g3doc/efficientnet-edgetpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/g3doc/efficientnet-edgetpu.png -------------------------------------------------------------------------------- /efficientnet_tf/tpu/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/tpu/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/__pycache__/eval_ckpt_main.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/__pycache__/eval_ckpt_main.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/__pycache__/lars_optimizer.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/__pycache__/lars_optimizer.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/__pycache__/preprocessing.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/__pycache__/preprocessing.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/lite/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/lite/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/condconv/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/condconv/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/edgetpu/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/edgetpu/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/__pycache__/efficientnet_builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/__pycache__/efficientnet_builder.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/__pycache__/efficientnet_model.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/__pycache__/efficientnet_model.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/__pycache__/model_builder_factory.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/__pycache__/model_builder_factory.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/condconv/__pycache__/condconv_layers.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/condconv/__pycache__/condconv_layers.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/tpu/__pycache__/efficientnet_x_builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/tpu/__pycache__/efficientnet_x_builder.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/lite/__pycache__/efficientnet_lite_builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/lite/__pycache__/efficientnet_lite_builder.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/edgetpu/__pycache__/efficientnet_edgetpu_builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/edgetpu/__pycache__/efficientnet_edgetpu_builder.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/condconv/__pycache__/efficientnet_condconv_builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xhluca/keras-noisy-student/HEAD/efficientnet_tf/condconv/__pycache__/efficientnet_condconv_builder.cpython-37.pyc -------------------------------------------------------------------------------- /efficientnet_tf/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /efficientnet_tf/lite/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /efficientnet_tf/tpu/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /efficientnet_tf/condconv/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /efficientnet_tf/edgetpu/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /temp.sh: -------------------------------------------------------------------------------- 1 | # export NAME="noisy_student_efficientnet-l2" 2 | # export MODEL="efficientnet-l2" 3 | # export URL="https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-l2.tar.gz" 4 | 5 | export NAME="noisy_student_efficientnet-b0" 6 | export MODEL="efficientnet-b0" 7 | export URL="https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b0.tar.gz" 8 | 9 | 10 | wget $URL -P tf_weights 11 | mkdir -p converted_weights 12 | mkdir -p tf_weights 13 | 14 | tar xvf "tf_weights/${NAME}.tar.gz" -C tf_weights 15 | 16 | python load_efficientnet.py \ 17 | --source ./efficientnet_tf \ 18 | --model_name $MODEL \ 19 | --tf_checkpoint "./tf_weights/${NAME}" \ 20 | --output_file "./converted_weights/${NAME}" -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | absl-py==0.11.0 2 | astunparse==1.6.3 3 | cachetools==4.2.0 4 | certifi==2020.12.5 5 | chardet==4.0.0 6 | cycler==0.10.0 7 | decorator==4.4.2 8 | efficientnet==1.1.1 9 | flatbuffers==1.12 10 | gast==0.3.3 11 | google-auth==1.24.0 12 | google-auth-oauthlib==0.4.2 13 | google-pasta==0.2.0 14 | grpcio==1.32.0 15 | h5py==2.10.0 16 | idna==2.10 17 | imageio==2.9.0 18 | importlib-metadata==3.3.0 19 | Keras-Applications==1.0.8 20 | Keras-Preprocessing==1.1.2 21 | kiwisolver==1.3.1 22 | Markdown==3.3.3 23 | matplotlib==3.3.3 24 | networkx==2.5 25 | numpy==1.19.4 26 | oauthlib==3.1.0 27 | opt-einsum==3.3.0 28 | Pillow==8.0.1 29 | protobuf==3.14.0 30 | pyasn1==0.4.8 31 | pyasn1-modules==0.2.8 32 | pyparsing==2.4.7 33 | python-dateutil==2.8.1 34 | PyWavelets==1.1.1 35 | requests==2.25.1 36 | requests-oauthlib==1.3.0 37 | rsa==4.6 38 | scikit-image==0.18.0 39 | scipy==1.5.4 40 | six==1.15.0 41 | tensorboard==2.4.0 42 | tensorboard-plugin-wit==1.7.0 43 | tensorflow==2.4.0 44 | tensorflow-estimator==2.4.0 45 | termcolor==1.1.0 46 | tifffile==2020.12.8 47 | typing-extensions==3.7.4.3 48 | urllib3==1.26.2 49 | Werkzeug==1.0.1 50 | wrapt==1.12.1 51 | zipp==3.4.0 52 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /efficientnet-b0 2 | /efficientnet-b0.tar.gz 3 | /temp 4 | tf_weights 5 | converted_weights 6 | efficientnet_tf/__pycache__ 7 | efficientnet_tf/__pycache__/efficientnet_builder.cpython-37.pyc 8 | efficientnet_tf/__pycache__/efficientnet_model.cpython-37.pyc 9 | efficientnet_tf/__pycache__/eval_ckpt_main.cpython-37.pyc 10 | efficientnet_tf/__pycache__/lars_optimizer.cpython-37.pyc 11 | efficientnet_tf/__pycache__/model_builder_factory.cpython-37.pyc 12 | efficientnet_tf/__pycache__/preprocessing.cpython-37.pyc 13 | efficientnet_tf/__pycache__/utils.cpython-37.pyc 14 | efficientnet_tf/condconv/__pycache__/__init__.cpython-37.pyc 15 | efficientnet_tf/condconv/__pycache__/condconv_layers.cpython-37.pyc 16 | efficientnet_tf/condconv/__pycache__/efficientnet_condconv_builder.cpython-37.pyc 17 | efficientnet_tf/edgetpu/__pycache__/__init__.cpython-37.pyc 18 | efficientnet_tf/edgetpu/__pycache__/efficientnet_edgetpu_builder.cpython-37.pyc 19 | efficientnet_tf/lite/__pycache__/__init__.cpython-37.pyc 20 | efficientnet_tf/lite/__pycache__/efficientnet_lite_builder.cpython-37.pyc 21 | efficientnet_tf/tpu/__pycache__/__init__.cpython-37.pyc 22 | efficientnet_tf/tpu/__pycache__/efficientnet_x_builder.cpython-37.pyc 23 | .vscode/settings.json 24 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Acknowledgment 2 | Code modified/retrieved from: 3 | 4 | * https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet 5 | * https://github.com/qubvel/efficientnet 6 | 7 | Please see respective license for more details. 8 | 9 | ## Download 10 | 11 | Here's how you can get the weights: 12 | * [Kaggle](https://www.kaggle.com/xhlulu/efn-l2) 13 | * [Github release](https://github.com/xhlulu/keras-efficientnet-l2/releases/tag/data) 14 | 15 | 16 | ## Instructions 17 | First, make sure to have the library and download the weights: 18 | ``` 19 | pip install efficientnet 20 | wget https://github.com/xhlulu/keras-efficientnet-l2/releases/download/data/efficientnet-l2_noisy-student.h5 21 | wget https://github.com/xhlulu/keras-efficientnet-l2/releases/download/data/efficientnet-l2_noisy-student_notop.h5 22 | ``` 23 | 24 | For `tensorflow>=2.4.0`: 25 | ```python 26 | import efficientnet.keras as efn 27 | 28 | model = efn.EfficientNetL2(weights="./efficientnet-l2_noisy-student_notop.h5", include_top=False) 29 | # or 30 | model = efn.EfficientNetL2(weights="./efficientnet-l2_noisy-student.h5", include_top=True) 31 | ``` 32 | 33 | For `tensorflow<=2.3.1`, there's a bug that would cause the L2 model to not load correctly. To use it, apply the following hack: 34 | ```python 35 | model = efn.EfficientNetL2( 36 | weights="./efficientnet-l2_noisy-student_notop.h5", 37 | include_top=False, 38 | drop_connect_rate=0 # the hack 39 | ) 40 | ``` 41 | 42 | However, this will modify the behavior of the model so you will need to be careful when using this. 43 | -------------------------------------------------------------------------------- /efficientnet_tf/lite/efficientnet_lite_builder_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for efficientnet_lite_builder.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import numpy as np 22 | import tensorflow.compat.v1 as tf 23 | 24 | from lite import efficientnet_lite_builder 25 | 26 | 27 | class EfficientnetBuilderTest(tf.test.TestCase): 28 | 29 | def _test_model_params(self, 30 | model_name, 31 | input_size, 32 | expected_params, 33 | override_params=None, 34 | features_only=False, 35 | pooled_features_only=False): 36 | images = tf.zeros((1, input_size, input_size, 3), dtype=tf.float32) 37 | efficientnet_lite_builder.build_model( 38 | images, 39 | model_name=model_name, 40 | override_params=override_params, 41 | training=True, 42 | features_only=features_only, 43 | pooled_features_only=pooled_features_only) 44 | num_params = np.sum([np.prod(v.shape) for v in tf.trainable_variables()]) 45 | 46 | self.assertEqual(num_params, expected_params) 47 | 48 | def test_efficientnet_b0(self): 49 | self._test_model_params( 50 | 'efficientnet-lite0', 224, expected_params=4652008) 51 | 52 | def test_efficientnet_b1(self): 53 | self._test_model_params( 54 | 'efficientnet-lite1', 240, expected_params=5416680) 55 | 56 | def test_efficientnet_b2(self): 57 | self._test_model_params( 58 | 'efficientnet-lite2', 260, expected_params=6092072) 59 | 60 | def test_efficientnet_b3(self): 61 | self._test_model_params( 62 | 'efficientnet-lite3', 280, expected_params=8197096) 63 | 64 | def test_efficientnet_b4(self): 65 | self._test_model_params( 66 | 'efficientnet-lite4', 300, expected_params=13006568) 67 | 68 | 69 | if __name__ == '__main__': 70 | tf.test.main() 71 | -------------------------------------------------------------------------------- /get_efn_weights.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | # ============================================================================= 3 | # Copyright 2019 Pavel Yakubovskiy, Sasha Illarionov. All Rights Reserved. 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # https://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================= 16 | 17 | import argparse 18 | import sys 19 | 20 | import numpy as np 21 | 22 | import tensorflow.compat.v1 as tf 23 | import efficientnet.tfkeras 24 | from tensorflow.keras.layers import BatchNormalization, Conv2D, Dense 25 | from efficientnet_tf import eval_ckpt_main 26 | from tqdm.auto import tqdm 27 | 28 | # def load_weights(model, weights): 29 | # """Load weights to Conv2D, BatchNorm, Dense layers of model sequentially""" 30 | # layer_index = 0 31 | # groupped_weights = group_weights(weights) 32 | # for layer in model.layers: 33 | # if isinstance(layer, (Conv2D, BatchNormalization, Dense)): 34 | # print(layer) 35 | # layer.set_weights(groupped_weights[layer_index]) 36 | # layer_index += 1 37 | 38 | 39 | model_name = "efficientnet-b0" 40 | model_ckpt = "./tf_weights/noisy_student_efficientnet-b0" 41 | output_file = "" 42 | example_img = "misc/panda.jpg" 43 | weights_only = True 44 | 45 | 46 | image_files = [example_img] 47 | eval_ckpt_driver = eval_ckpt_main.EvalCkptDriver(model_name) 48 | with tf.Graph().as_default(), tf.Session() as sess: 49 | images, _ = eval_ckpt_driver.build_dataset( 50 | image_files, [0] * len(image_files), False 51 | ) 52 | eval_ckpt_driver.build_model(images, is_training=False) 53 | sess.run(tf.global_variables_initializer()) 54 | eval_ckpt_driver.restore_model(sess, model_ckpt) 55 | global_variables = tf.global_variables() 56 | weights = dict() 57 | print("Starting!") 58 | for variable in tqdm(global_variables): 59 | try: 60 | weights[variable.name] = variable.eval() 61 | except: 62 | print(f"Skipping variable {variable.name}, an exception occurred") 63 | # model = _get_model_by_name( 64 | # model_name, include_top=True, input_shape=None, weights=None, classes=1000 65 | # ) 66 | # load_weights(model, weights) 67 | # output_file = f"{output_file}.h5" 68 | # if weights_only: 69 | # model.save_weights(output_file) 70 | # else: 71 | # model.save(output_file) 72 | -------------------------------------------------------------------------------- /efficientnet_tf/model_builder_factory.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Utilities for model builder or input size.""" 16 | 17 | from . import efficientnet_builder 18 | from .condconv import efficientnet_condconv_builder 19 | from .edgetpu import efficientnet_edgetpu_builder 20 | from .lite import efficientnet_lite_builder 21 | from .tpu import efficientnet_x_builder 22 | 23 | 24 | def get_model_builder(model_name): 25 | """Get the model_builder module for a given model name.""" 26 | if model_name.startswith('efficientnet-lite'): 27 | return efficientnet_lite_builder 28 | elif model_name.startswith('efficientnet-edgetpu-'): 29 | return efficientnet_edgetpu_builder 30 | elif model_name.startswith('efficientnet-condconv-'): 31 | return efficientnet_condconv_builder 32 | elif model_name.startswith('efficientnet-x-'): 33 | return efficientnet_x_builder 34 | elif model_name.startswith('efficientnet-'): 35 | return efficientnet_builder 36 | else: 37 | raise ValueError( 38 | 'Model must be either efficientnet-b* or efficientnet-edgetpu* or' 39 | 'efficientnet-condconv*, efficientnet-lite*') 40 | 41 | 42 | def get_model_input_size(model_name): 43 | """Get model input size for a given model name.""" 44 | if model_name.startswith('efficientnet-lite'): 45 | _, _, image_size, _ = ( 46 | efficientnet_lite_builder.efficientnet_lite_params(model_name)) 47 | elif model_name.startswith('efficientnet-edgetpu-'): 48 | _, _, image_size, _ = ( 49 | efficientnet_edgetpu_builder.efficientnet_edgetpu_params(model_name)) 50 | elif model_name.startswith('efficientnet-condconv-'): 51 | _, _, image_size, _, _ = ( 52 | efficientnet_condconv_builder.efficientnet_condconv_params(model_name)) 53 | elif model_name.startswith('efficientnet-x'): 54 | _, _, image_size, _, _ = efficientnet_x_builder.efficientnet_x_params( 55 | model_name) 56 | elif model_name.startswith('efficientnet'): 57 | _, _, image_size, _ = efficientnet_builder.efficientnet_params(model_name) 58 | else: 59 | raise ValueError( 60 | 'Model must be either efficientnet-b* or efficientnet-x-b* or efficientnet-edgetpu* or ' 61 | 'efficientnet-condconv*, efficientnet-lite*') 62 | return image_size 63 | 64 | -------------------------------------------------------------------------------- /efficientnet_tf/inspect_model_architecture.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Export a dummy-quantized tflite model corresponding to the given model.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | from absl import app 22 | from absl import flags 23 | import tensorflow as tf 24 | 25 | import efficientnet_builder 26 | from edgetpu import efficientnet_edgetpu_builder 27 | 28 | flags.DEFINE_string('model_name', 'efficientnet-b0', 'Model name to inspect.') 29 | flags.DEFINE_integer('image_res', 224, 'The size of the input image') 30 | flags.DEFINE_string('output_tflite', '/tmp/model.tflite', 31 | 'Location of the generated tflite model') 32 | 33 | # FLAGS should not be used before main. 34 | FLAGS = flags.FLAGS 35 | 36 | 37 | def main(unused_argv): 38 | tf.logging.set_verbosity(tf.logging.ERROR) 39 | image_res = FLAGS.image_res 40 | model_name = FLAGS.model_name 41 | model_builder_fn = None 42 | if model_name.startswith('efficientnet-edgetpu'): 43 | model_builder_fn = efficientnet_edgetpu_builder 44 | elif model_name.startswith('efficientnet'): 45 | model_builder_fn = efficientnet_builder 46 | 47 | else: 48 | raise ValueError( 49 | 'Model must be either efficientnet-b* or efficientnet-edgetpu*') 50 | 51 | with tf.Graph().as_default(), tf.Session() as sess: 52 | images = tf.placeholder( 53 | tf.float32, shape=(1, image_res, image_res, 3), name='input') 54 | output, _ = model_builder_fn.build_model( 55 | images, FLAGS.model_name, training=False) 56 | 57 | tf.global_variables_initializer().run() 58 | updates = [] 59 | for var in tf.trainable_variables(): 60 | noise = tf.random.normal(shape=var.shape, stddev=0.001) 61 | updates.append(var.assign_add(noise)) 62 | sess.run(updates) 63 | converter = tf.lite.TFLiteConverter.from_session(sess, [images], [output]) # pytype: disable=attribute-error 64 | converter.inference_type = tf.lite.constants.QUANTIZED_UINT8 65 | converter.quantized_input_stats = {'input': (0, 1.)} 66 | converter.default_ranges_stats = (-10, 10) 67 | 68 | tflite_model = converter.convert() 69 | tf.gfile.Open(FLAGS.output_tflite, 'wb').write(tflite_model) 70 | 71 | print('Model %s, image size %d' % (model_name, image_res)) 72 | print('TfLite model stored at %s' % FLAGS.output_tflite) 73 | 74 | 75 | if __name__ == '__main__': 76 | app.run(main) 77 | -------------------------------------------------------------------------------- /efficientnet_tf/condconv/README.md: -------------------------------------------------------------------------------- 1 | # EfficientNet-CondConv 2 | 3 | [1] Brandon Yang, Gabriel Bender, Quoc V. Le, Jiquan Ngiam. CondConv: Conditionally Parameterized Convolutions for Efficient Inference. NeurIPS 2019. Arxiv Link: https://arxiv.org/abs/1904.04971. 4 | 5 | ## 1. About CondConv 6 | 7 | Conditionally parameterized convolutions (CondConv) are a new building block for convolutional neural networks to increase capacity while maintaining efficient inference. In a traditional convolutional layer, each example is processed with the same kernel. In a CondConv layer, each example is processed with a specialized, example-dependent kernel. As an intuitive motivating example, on the ImageNet classification dataset, we might want to classify dogs and cats with different convolutional kernels. 8 | 9 | 10 | 11 | 14 | 15 |
12 | 13 |
16 | 17 | A CondConv layer consists of n experts, each of which are the same size as the convolutional kernel of the original convolutional layer. For each example, the example-dependent convolutional kernel is computed as the weighted sum of experts using an example-dependent routing function. Increasing the number of experts enables us to increase the capacity of a network, while maintaining efficient inference. 18 | 19 | Replacing convolutional layers with CondConv layers improves the accuracy versus inference cost trade-off on a wide range of models: MobileNetV1, MobileNetV2, ResNets, and EfficientNets. We measure inference cost in multiply-adds (MADDs). When applied to EfficientNets, we obtain EfficientNet-CondConv models. Our EfficientNet-CondConv-B0 model with 8 experts achieves state-of-the-art accuracy versus inference cost performance. 20 | 21 | In this directory, we open-source the code to reproduce the EfficientNet-CondConv results in our paper and enable easy experimentation with EfficientNet-CondConv models. Additionally, we open-source the CondConv2d and DepthwiseCondConv2D Keras layers for easy application in new model architectures. 22 | 23 | ## 2. Using pretrained EfficientNet-CondConv checkpoints 24 | 25 | We have provided pre-trained checkpoints for several EfficientNet-CondConv models. 26 | 27 | | | CondConv Experts | Params | MADDs | Accuracy | 28 | |--------------------------------|------------------|--------|-------|----------| 29 | | EfficientNet-B0 | - | 5.3M | 391M | 77.3 | 30 | | EfficientNet-CondConv-B0 ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/condconv/efficientnet-condconv-b0-4e.tar.gz))| 4 | 13.3M | 402M | 77.8 | 31 | | EfficientNet-CondConv-B0 ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/condconv/efficientnet-condconv-b0-8e.tar.gz))| 8 | 24.0M | 413M | 78.3 | 32 | 33 | | | CondConv Experts | Params | MADDs | Accuracy | 34 | |---------------------------------------|------------------|--------|-------|----------| 35 | | EfficientNet-B1 | - | 7.8M | 700M | 79.2 | 36 | | EfficientNet-CondConv-B0-Depth ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/condconv/efficientnet-condconv-b0-8e-depth.tar.gz)) | 8 | 39.7M | 614M | 79.5 | 37 | 38 | A quick way to use these checkpoints is to run: 39 | 40 | ```shell 41 | $ export MODEL=efficientnet-condconv-b0-8e 42 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz 43 | $ tar zxf ${MODEL}.tar.gz 44 | $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg 45 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.txt 46 | $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.txt 47 | ``` 48 | 49 | Please refer to the following colab for more instructions on how to obtain and use those checkpoints. 50 | 51 | * [`eval_ckpt_example.ipynb`](eval_ckpt_example.ipynb): A colab example to load 52 | EfficientNet pretrained checkpoints files and use the restored model to classify images. 53 | 54 | ## 3. Training EfficientNet-CondConv models on Cloud TPUs 55 | Please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/efficientnet. 56 | 57 | -------------------------------------------------------------------------------- /efficientnet_tf/edgetpu/README.md: -------------------------------------------------------------------------------- 1 | # EfficientNet-EdgeTPU 2 | 3 | **Blog post: https://ai.googleblog.com/2019/08/efficientnet-edgetpu-creating.html** 4 | 5 | EfficientNet-EdgeTPU are a family of image classification neural network models customized for deployment on [Google Edge TPU](https://coral.withgoogle.com/). These networks are closely related to [EfficientNets] (https://arxiv.org/abs/1905.11946). 6 | 7 | EfficientNet-EdgeTPU were developed using the [AutoML MNAS framework](https://ai.googleblog.com/2018/08/mnasnet-towards-automating-design-of.html) by augmenting the neural network search space with building blocks tuned to execute efficiently on the EdgeTPU neural network accelerator architecture. The neural architecture search was incentivized to discover models that achieve low parameter footprint and low latency on EdgeTpu, while simultaneously achieving high classification accuracy. This neural architecture search produced a baseline model: edgetpunet-S, which is subsequently scaled up using EfficientNet's compound scaling method to produce the M and L models. 8 | 9 | 10 | 11 | 14 | 15 |
12 | 13 |
16 | 17 | ### Using Pretrained EfficientNet-EdgeTPU Checkpoints 18 | We have provided pretrained checkpoints and float/quantized TFLite models: 19 | 20 | * [EfficientNet-EdgeTPU-S](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/efficientnet-edgetpu-S.tar.gz) 21 | * [EfficientNet-EdgeTPU-M](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/efficientnet-edgetpu-M.tar.gz) 22 | * [EfficientNet-EdgeTPU-L](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/efficientnet-edgetpu-L.tar.gz) 23 | 24 | A quick way to use these checkpoints is to run: 25 | 26 | ```shell 27 | $ export MODEL=efficientnet-edgetpu-S 28 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz 29 | $ tar zxf ${MODEL}.tar.gz 30 | $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg 31 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.txt 32 | $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.txt --include_background_label 33 | ``` 34 | 35 | Note that these models were trained with label#0 marked as the background label for easier deployment. TFLite models can be evaluated using this [tool](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/evaluation/tasks/imagenet_image_classification). 36 | 37 | ### Training EfficientNet-EdgeTPU on Cloud TPUs 38 | Please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/efficientnet 39 | 40 | ### Post-training quantization 41 | 42 | EdgeTPUs support inference using integer quantized models only. We found that using the [Tensorflow Lite's post-training quantization tool](https://www.tensorflow.org/lite/performance/post_training_quantization) works remarkably well for producing a EdgeTPU-compatible quantized model from a floating-point training checkpoint. For full integer quantization, the post-training quantization tool requires a representative dataset for calibrating the dynamic ranges of the activations. 43 | 44 | We provide a tool that invokes the post-training quantization tool to produce quantized tensorflow-lite model: 45 | 46 | ```shell 47 | $ export MODEL=efficientnet-edgetpu-S 48 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz 49 | $ tar zxf ${MODEL}.tar.gz 50 | $ python export_model.py --model_name=$MODEL --ckpt_dir=$MODEL --data_dir=/path/to/representative_dataset/ --output_tflite=${MODEL}_quant.tflite 51 | ``` 52 | 53 | To produce a float model that bypasses the post-training quantization: 54 | 55 | ```shell 56 | $ python export_model.py --model_name=$MODEL --ckpt_dir=$MODEL --output_tflite=${MODEL}_float.tflite --quantize=False 57 | ``` 58 | The table below compared the accuracy of float models (on CPU) and the quantized models on EdgeTPU: 59 | 60 | |**Model** | **Imagenet top-1 accuracy (float)** | **Imagenet top-1 accuracy (quantized)** | 61 | |------|----------------------|------------------| 62 | |efficientnet-edgetpu-S| 77.23% | 77.0 % | 63 | |efficientnet-edgetpu-M| 78.69 | 78.6 % | 64 | |efficientnet-edgetpu-L| 80.62 | 80.2% | 65 | 66 | The `export_model.py` script can also be used to export a [tensorflow saved_model](https://www.tensorflow.org/guide/saved_model) from a training checkpoint: 67 | 68 | ```shell 69 | $ python export_model.py --model_name=$MODEL --ckpt_dir=/path/to/model-ckpt/ --output_saved_model_dir=/path/to/output_saved_model/ --output_tflite=${MODEL}_float.tflite --quantize=False 70 | ``` 71 | -------------------------------------------------------------------------------- /efficientnet_tf/lite/README.md: -------------------------------------------------------------------------------- 1 | # EfficientNet-lite 2 | 3 | EfficientNet-lite are a set of mobile/IoT friendly image classification models. Notably, while EfficientNet-EdgeTPU that is specialized for Coral EdgeTPU, these EfficientNet-lite models run well on all mobile CPU/GPU/EdgeTPU. 4 | 5 | Due to the requirements from edge devices, we mainly made the following changes based on the original EfficientNets. 6 | 7 | * Remove squeeze-and-excite (SE): SE are not well supported for some mobile accelerators. 8 | * Replace all swish with RELU6: for easier post-quantization. 9 | * Fix the stem and head while scaling models up: for keeping models small and fast. 10 | 11 | 12 | Here are the checkpoints, and their accurracy, params, flops, and Pixel4's CPU/GPU/EdgeTPU latency. 13 | 14 | |**Model** | **params** | **MAdds** | **FP32 accuracy** | **FP32 CPU latency** | **FP32 GPU latency** | **FP16 GPU latency** |**INT8 accuracy** | **INT8 CPU latency** | **INT8 TPU latency**| 15 | |------|-----|-------|-------|-------|-------|-------|-------|-------|-------| 16 | |efficientnet-lite0 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite0.tar.gz) | 4.7M | 407M | 75.1% | 12ms | 9.0ms | 6.0ms | 74.4% | 6.5ms | 3.8ms | 17 | |efficientnet-lite1 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite1.tar.gz) | 5.4M | 631M | 76.7% | 18ms | 12ms | 8.0ms | 75.9% | 9.1ms | 5.4ms | 18 | |efficientnet-lite2 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite2.tar.gz) | 6.1M | 899M | 77.6% | 26ms | 16ms | 10ms | 77.0% | 12ms | 7.9ms | 19 | |efficientnet-lite3 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite3.tar.gz) | 8.2M | 1.44B | 79.8% | 41ms | 23ms | 14ms | 79.0% | 18ms | 9.7ms | 20 | |efficientnet-lite4 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite4.tar.gz) |13.0M | 2.64B | 81.5% | 76ms | 36ms | 21ms | 80.2% | 30ms | - | 21 | 22 | * CPU/GPU/TPU latency are measured on Pixel4, with batch size 1 and 4 CPU threads. FP16 GPU latency is measured with default latency, while FP32 GPU latency is measured with additional option --gpu_precision_loss_allowed=false. 23 | 24 | * Each checkpoint all contains FP tflite and post-training quantized INT8 tflite files. If you use these models or checkpoints, you can cite this [efficientnet paper](https://arxiv.org/abs/1905.11946). 25 | 26 | Comparing with MobileNetV2, ResNet-50, and Inception-V4, our models have 27 | better trade-offs between accuracy and size/latency. 28 | The following two figures show the comparison among quantized versions of 29 | these models. The latency numbers are obtained on a Pixel 4 with 4 CPU 30 | threads. 31 | 32 |

33 | 34 |

35 | 36 |

37 | 38 |

39 | 40 | As Tensorflow Lite also provides GPU acceleration for float models, the 41 | following shows the latency comparison among float versions of these 42 | models. Again, the latency numbers are obtained on a Pixel 4. 43 |

44 | 45 |

46 | 47 | A quick way to use these checkpoints is to run: 48 | 49 | ```shell 50 | $ export MODEL=efficientnet-lite0 51 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz 52 | $ tar zxf ${MODEL}.tar.gz 53 | $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg 54 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.txt 55 | $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.txt 56 | ``` 57 | 58 | TFLite models can be evaluated using this [tool](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/evaluation/tasks/imagenet_image_classification). 59 | 60 | ### Training EfficientNet-lite on Cloud TPUs 61 | Please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/efficientnet 62 | 63 | ### Post-training quantization 64 | 65 | 66 | ```shell 67 | $ export MODEL=efficientnet-lite0 68 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz 69 | $ tar zxf ${MODEL}.tar.gz 70 | $ python export_model.py --model_name=$MODEL --ckpt_dir=$MODEL --data_dir=/path/to/representative_dataset/ --output_tflite=${MODEL}_quant.tflite 71 | ``` 72 | 73 | To produce a float model that bypasses the post-training quantization: 74 | 75 | ```shell 76 | $ python export_model.py --model_name=$MODEL --ckpt_dir=$MODEL --output_tflite=${MODEL}_float.tflite --quantize=False 77 | ``` 78 | 79 | The `export_model.py` script can also be used to export a [tensorflow saved_model](https://www.tensorflow.org/guide/saved_model) from a training checkpoint: 80 | 81 | ```shell 82 | $ python export_model.py --model_name=$MODEL --ckpt_dir=/path/to/model-ckpt/ --output_saved_model_dir=/path/to/output_saved_model/ --output_tflite=${MODEL}_float.tflite --quantize=False 83 | ``` 84 | -------------------------------------------------------------------------------- /efficientnet_tf/eval_ckpt_main.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Eval checkpoint driver. 16 | 17 | This is an example evaluation script for users to understand the EfficientNet 18 | model checkpoints on CPU. To serve EfficientNet, please consider to export a 19 | `SavedModel` from checkpoints and use tf-serving to serve. 20 | """ 21 | 22 | from __future__ import absolute_import 23 | from __future__ import division 24 | from __future__ import print_function 25 | 26 | from absl import app 27 | from absl import flags 28 | from absl import logging 29 | import tensorflow.compat.v1 as tf 30 | 31 | from . import model_builder_factory 32 | from . import preprocessing 33 | from . import utils 34 | 35 | flags.DEFINE_string('model_name', 'efficientnet-b0', 'Model name to eval.') 36 | flags.DEFINE_string('runmode', 'examples', 'Running mode: examples or imagenet') 37 | flags.DEFINE_string( 38 | 'imagenet_eval_glob', None, 'Imagenet eval image glob, ' 39 | 'such as /imagenet/ILSVRC2012*.JPEG') 40 | flags.DEFINE_string( 41 | 'imagenet_eval_label', None, 'Imagenet eval label file path, ' 42 | 'such as /imagenet/ILSVRC2012_validation_ground_truth.txt') 43 | flags.DEFINE_string('ckpt_dir', '/tmp/ckpt/', 'Checkpoint folders') 44 | flags.DEFINE_boolean('enable_ema', True, 'Enable exponential moving average.') 45 | flags.DEFINE_string('export_ckpt', None, 'Exported ckpt for eval graph.') 46 | flags.DEFINE_string('example_img', '/tmp/panda.jpg', 47 | 'Filepath for a single example image.') 48 | flags.DEFINE_string('labels_map_file', '/tmp/labels_map.txt', 49 | 'Labels map from label id to its meaning.') 50 | flags.DEFINE_bool('include_background_label', False, 51 | 'Whether to include background as label #0') 52 | flags.DEFINE_bool('advprop_preprocessing', False, 53 | 'Whether to use AdvProp preprocessing.') 54 | flags.DEFINE_integer('num_images', 5000, 55 | 'Number of images to eval. Use -1 to eval all images.') 56 | 57 | 58 | class EvalCkptDriver(utils.EvalCkptDriver): 59 | """A driver for running eval inference.""" 60 | 61 | def build_model(self, features, is_training): 62 | """Build model with input features.""" 63 | tf.logging.info(self.model_name) 64 | model_builder = model_builder_factory.get_model_builder(self.model_name) 65 | 66 | if self.advprop_preprocessing: 67 | # AdvProp uses Inception preprocessing. 68 | features = features * 2.0 / 255 - 1.0 69 | else: 70 | features -= tf.constant( 71 | model_builder.MEAN_RGB, shape=[1, 1, 3], dtype=features.dtype) 72 | features /= tf.constant( 73 | model_builder.STDDEV_RGB, shape=[1, 1, 3], dtype=features.dtype) 74 | logits, _ = model_builder.build_model( 75 | features, self.model_name, is_training) 76 | probs = tf.nn.softmax(logits) 77 | probs = tf.squeeze(probs) 78 | return probs 79 | 80 | def get_preprocess_fn(self): 81 | """Build input dataset.""" 82 | return preprocessing.preprocess_image 83 | 84 | 85 | def get_eval_driver(model_name, 86 | include_background_label=False, 87 | advprop_preprocessing=False): 88 | """Get a eval driver.""" 89 | image_size = model_builder_factory.get_model_input_size(model_name) 90 | return EvalCkptDriver( 91 | model_name=model_name, 92 | batch_size=1, 93 | image_size=image_size, 94 | include_background_label=include_background_label, 95 | advprop_preprocessing=advprop_preprocessing) 96 | 97 | 98 | # FLAGS should not be used before main. 99 | FLAGS = flags.FLAGS 100 | 101 | 102 | def main(unused_argv): 103 | logging.set_verbosity(logging.ERROR) 104 | driver = get_eval_driver(FLAGS.model_name, FLAGS.include_background_label, 105 | FLAGS.advprop_preprocessing) 106 | if FLAGS.runmode == 'examples': 107 | # Run inference for an example image. 108 | driver.eval_example_images(FLAGS.ckpt_dir, [FLAGS.example_img], 109 | FLAGS.labels_map_file, FLAGS.enable_ema, 110 | FLAGS.export_ckpt) 111 | elif FLAGS.runmode == 'imagenet': 112 | # Run inference for imagenet. 113 | driver.eval_imagenet(FLAGS.ckpt_dir, FLAGS.imagenet_eval_glob, 114 | FLAGS.imagenet_eval_label, FLAGS.num_images, 115 | FLAGS.enable_ema, FLAGS.export_ckpt) 116 | else: 117 | print('must specify runmode: examples or imagenet') 118 | 119 | 120 | if __name__ == '__main__': 121 | app.run(main) 122 | -------------------------------------------------------------------------------- /load_efficientnet.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2019 Pavel Yakubovskiy, Sasha Illarionov. All Rights Reserved. 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================= 15 | 16 | import argparse 17 | import sys 18 | 19 | import numpy as np 20 | 21 | import tensorflow.compat.v1 as tf 22 | # import efficientnet.tfkeras 23 | from tensorflow.python.keras.applications import efficientnet as efn 24 | from tensorflow.keras.layers import BatchNormalization, Conv2D, Dense 25 | 26 | tf.logging.set_verbosity("INFO") 27 | 28 | 29 | def EfficientNetL2( 30 | include_top=True, 31 | weights='imagenet', 32 | input_tensor=None, 33 | input_shape=None, 34 | pooling=None, 35 | classes=1000, 36 | **kwargs 37 | ): 38 | return efn.EfficientNet( 39 | 4.3, 5.3, 800, 0.5, 40 | model_name='efficientnet-l2', 41 | include_top=include_top, weights=weights, 42 | input_tensor=input_tensor, input_shape=input_shape, 43 | pooling=pooling, classes=classes, 44 | **kwargs 45 | ) 46 | 47 | 48 | def _get_model_by_name(name, *args, **kwargs): 49 | models = { 50 | 'efficientnet-b0': efn.EfficientNetB0, 51 | 'efficientnet-b1': efn.EfficientNetB1, 52 | 'efficientnet-b2': efn.EfficientNetB2, 53 | 'efficientnet-b3': efn.EfficientNetB3, 54 | 'efficientnet-b4': efn.EfficientNetB4, 55 | 'efficientnet-b5': efn.EfficientNetB5, 56 | 'efficientnet-b6': efn.EfficientNetB6, 57 | 'efficientnet-b7': efn.EfficientNetB7, 58 | 'efficientnet-l2': EfficientNetL2, 59 | } 60 | 61 | model_fn = models[name] 62 | model = model_fn(*args, **kwargs) 63 | return model 64 | 65 | 66 | def group_weights(weights): 67 | """ 68 | Group each layer weights together, initially all weights are dict of 'layer_name/layer_var': np.array 69 | 70 | Example: 71 | input: { 72 | ...: ... 73 | 'conv2d/kernel': , 74 | 'conv2d/bias': , 75 | ...: ... 76 | } 77 | output: [..., [...], [, ], [...], ...] 78 | 79 | """ 80 | 81 | out_weights = [] 82 | 83 | previous_layer_name = "" 84 | group = [] 85 | 86 | for k, v in weights.items(): 87 | 88 | layer_name = "/".join(k.split("/")[:-1]) 89 | 90 | if layer_name == previous_layer_name: 91 | group.append(v) 92 | else: 93 | if group: 94 | out_weights.append(group) 95 | 96 | group = [v] 97 | previous_layer_name = layer_name 98 | 99 | out_weights.append(group) 100 | return out_weights 101 | 102 | 103 | def load_weights(model, weights): 104 | """Load weights to Conv2D, BatchNorm, Dense layers of model sequentially""" 105 | layer_index = 0 106 | groupped_weights = group_weights(weights) 107 | for layer in model.layers: 108 | if isinstance(layer, (Conv2D, BatchNormalization, Dense)): 109 | print(layer) 110 | layer.set_weights(groupped_weights[layer_index]) 111 | layer_index += 1 112 | 113 | 114 | def convert_tensorflow_model( 115 | model, model_name, model_ckpt, output_file, example_img="misc/panda.jpg", weights_only=True 116 | ): 117 | """ Loads and saves a TensorFlow model. """ 118 | image_files = [example_img] 119 | eval_ckpt_driver = eval_ckpt_main.EvalCkptDriver(model_name) 120 | with tf.Graph().as_default(), tf.Session() as sess: 121 | images, _ = eval_ckpt_driver.build_dataset( 122 | image_files, [0] * len(image_files), False 123 | ) 124 | eval_ckpt_driver.build_model(images, is_training=False) 125 | sess.run(tf.global_variables_initializer()) 126 | eval_ckpt_driver.restore_model(sess, model_ckpt) 127 | global_variables = tf.global_variables() 128 | weights = dict() 129 | for variable in global_variables: 130 | try: 131 | weights[variable.name] = variable.eval() 132 | except: 133 | print(f"Skipping variable {variable.name}, an exception occurred") 134 | 135 | load_weights(model, weights) 136 | output_file = f"{output_file}.h5" 137 | if weights_only: 138 | model.save_weights(output_file) 139 | else: 140 | model.save(output_file) 141 | 142 | 143 | if __name__ == "__main__": 144 | parser = argparse.ArgumentParser( 145 | description="Convert TF model to Keras and save for easier future loading" 146 | ) 147 | parser.add_argument( 148 | "--source", type=str, default="dist/tf_src", help="source code path" 149 | ) 150 | parser.add_argument( 151 | "--model_name", 152 | type=str, 153 | default="efficientnet-b0", 154 | help="efficientnet-b{N}, where N is an integer 0 <= N <= 7", 155 | ) 156 | parser.add_argument( 157 | "--tf_checkpoint", 158 | type=str, 159 | default="pretrained_tensorflow/efficientnet-b0/", 160 | help="checkpoint file path", 161 | ) 162 | parser.add_argument( 163 | "--output_file", 164 | type=str, 165 | default="pretrained_keras/efficientnet-b0", 166 | help="output Keras model file name", 167 | ) 168 | parser.add_argument( 169 | "--weights_only", 170 | type=str, 171 | default="true", 172 | help="Whether to include metadata in the serialized Keras model", 173 | ) 174 | args = parser.parse_args() 175 | 176 | sys.path.append(args.source) 177 | from efficientnet_tf import eval_ckpt_main 178 | 179 | true_values = ("yes", "true", "t", "1", "y") 180 | 181 | model = _get_model_by_name( 182 | args.model_name, include_top=True, input_shape=None, weights=None, classes=1000) 183 | convert_tensorflow_model( 184 | model=model, 185 | model_name=args.model_name, 186 | model_ckpt=args.tf_checkpoint, 187 | output_file=args.output_file, 188 | weights_only=args.weights_only in true_values, 189 | ) 190 | -------------------------------------------------------------------------------- /efficientnet_tf/export_model.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Export model (float or quantized tflite, and saved model) from a trained checkpoint.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | from absl import app 22 | from absl import flags 23 | import tensorflow.compat.v1 as tf 24 | 25 | import imagenet_input 26 | import model_builder_factory 27 | 28 | flags.DEFINE_string("model_name", None, "Model name to eval.") 29 | flags.DEFINE_string("ckpt_dir", None, "Path to the training checkpoint") 30 | flags.DEFINE_boolean("enable_ema", True, "Enable exponential moving average.") 31 | flags.DEFINE_string("data_dir", None, 32 | "Image dataset directory for post training quantization.") 33 | flags.DEFINE_string("output_tflite", None, "Path to output tflite file.") 34 | flags.DEFINE_bool("quantize", True, 35 | "Quantize model to uint8 before exporting tflite model.") 36 | flags.DEFINE_integer( 37 | "num_steps", 2000, 38 | "Number of post-training quantization calibration steps to run.") 39 | flags.DEFINE_integer("image_size", 224, "Size of the input image.") 40 | flags.DEFINE_integer("batch_size", 1, "Batch size of input tensor.") 41 | flags.DEFINE_string("endpoint_name", None, "Endpoint name") 42 | flags.DEFINE_string("output_saved_model_dir", None, 43 | "Directory in which to save the saved_model.") 44 | FLAGS = flags.FLAGS 45 | 46 | 47 | def restore_model(sess, ckpt_dir, enable_ema=True): 48 | """Restore variables from checkpoint dir.""" 49 | sess.run(tf.global_variables_initializer()) 50 | checkpoint = tf.train.latest_checkpoint(ckpt_dir) 51 | if enable_ema: 52 | ema = tf.train.ExponentialMovingAverage(decay=0.0) 53 | ema_vars = tf.trainable_variables() + tf.get_collection("moving_vars") 54 | for v in tf.global_variables(): 55 | if "moving_mean" in v.name or "moving_variance" in v.name: 56 | ema_vars.append(v) 57 | ema_vars = list(set(ema_vars)) 58 | var_dict = ema.variables_to_restore(ema_vars) 59 | else: 60 | var_dict = None 61 | 62 | sess.run(tf.global_variables_initializer()) 63 | saver = tf.train.Saver(var_dict, max_to_keep=1) 64 | saver.restore(sess, checkpoint) 65 | 66 | 67 | def representative_dataset_gen(): 68 | """Gets a python generator of image numpy arrays for ImageNet.""" 69 | params = dict(batch_size=FLAGS.batch_size) 70 | imagenet_eval = imagenet_input.ImageNetInput( 71 | is_training=False, 72 | data_dir=FLAGS.data_dir, 73 | transpose_input=False, 74 | cache=False, 75 | image_size=FLAGS.image_size, 76 | num_parallel_calls=1, 77 | use_bfloat16=False, 78 | include_background_label=True, 79 | ) 80 | 81 | data = imagenet_eval.input_fn(params) 82 | 83 | def preprocess_map_fn(images, labels): 84 | del labels 85 | model_builder = model_builder_factory.get_model_builder(FLAGS.model_name) 86 | images -= tf.constant( 87 | model_builder.MEAN_RGB, shape=[1, 1, 3], dtype=images.dtype) 88 | images /= tf.constant( 89 | model_builder.STDDEV_RGB, shape=[1, 1, 3], dtype=images.dtype) 90 | return images 91 | 92 | data = data.map(preprocess_map_fn) 93 | iterator = data.make_one_shot_iterator() 94 | for _ in range(FLAGS.num_steps): 95 | # In eager context, we can get a python generator from a dataset iterator. 96 | images = iterator.get_next() 97 | yield [images] 98 | 99 | 100 | def main(_): 101 | # Enables eager context for TF 1.x. TF 2.x will use eager by default. 102 | # This is used to conveniently get a representative dataset generator using 103 | # TensorFlow training input helper. 104 | tf.enable_eager_execution() 105 | 106 | model_builder = model_builder_factory.get_model_builder(FLAGS.model_name) 107 | 108 | with tf.Graph().as_default(), tf.Session() as sess: 109 | images = tf.placeholder( 110 | tf.float32, 111 | shape=(1, FLAGS.image_size, FLAGS.image_size, 3), 112 | name="images") 113 | 114 | logits, endpoints = model_builder.build_model(images, FLAGS.model_name, 115 | False) 116 | if FLAGS.endpoint_name: 117 | output_tensor = endpoints[FLAGS.endpoint_name] 118 | else: 119 | output_tensor = tf.nn.softmax(logits) 120 | 121 | restore_model(sess, FLAGS.ckpt_dir, FLAGS.enable_ema) 122 | 123 | if FLAGS.output_saved_model_dir: 124 | signature_def_map = { 125 | "serving_default": 126 | tf.compat.v1.saved_model.signature_def_utils 127 | .predict_signature_def({"input": images}, 128 | {"output": output_tensor}) 129 | } 130 | 131 | builder = tf.compat.v1.saved_model.Builder(FLAGS.output_saved_model_dir) 132 | builder.add_meta_graph_and_variables( 133 | sess, ["serve"], signature_def_map=signature_def_map) 134 | builder.save() 135 | print("Saved model written to %s" % FLAGS.output_saved_model_dir) 136 | 137 | converter = tf.lite.TFLiteConverter.from_session(sess, [images], 138 | [output_tensor]) 139 | if FLAGS.quantize: 140 | if not FLAGS.data_dir: 141 | raise ValueError( 142 | "Post training quantization requires data_dir flag to point to the " 143 | "calibration dataset. To export a float model, set " 144 | "--quantize=False.") 145 | 146 | converter.representative_dataset = tf.lite.RepresentativeDataset( 147 | representative_dataset_gen) 148 | converter.optimizations = [tf.lite.Optimize.DEFAULT] 149 | converter.inference_input_type = tf.lite.constants.QUANTIZED_UINT8 150 | converter.inference_output_type = tf.lite.constants.QUANTIZED_UINT8 151 | converter.target_spec.supported_ops = [ 152 | tf.lite.OpsSet.TFLITE_BUILTINS_INT8 153 | ] 154 | 155 | tflite_buffer = converter.convert() 156 | tf.gfile.GFile(FLAGS.output_tflite, "wb").write(tflite_buffer) 157 | print("tflite model written to %s" % FLAGS.output_tflite) 158 | 159 | 160 | if __name__ == "__main__": 161 | flags.mark_flag_as_required("model_name") 162 | flags.mark_flag_as_required("ckpt_dir") 163 | flags.mark_flag_as_required("output_tflite") 164 | app.run(main) 165 | -------------------------------------------------------------------------------- /efficientnet_tf/lars_optimizer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Layer-wise Adaptive Rate Scaling optimizer for large-batch training.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import tensorflow.compat.v1 as tf 22 | 23 | 24 | class LARSOptimizer(tf.train.Optimizer): 25 | """Layer-wise Adaptive Rate Scaling for large batch training. 26 | 27 | Introduced by "Large Batch Training of Convolutional Networks" by Y. You, 28 | I. Gitman, and B. Ginsburg. (https://arxiv.org/abs/1708.03888) 29 | 30 | Implements the LARS learning rate scheme presented in the paper above. This 31 | optimizer is useful when scaling the batch size to up to 32K without 32 | significant performance degradation. It is recommended to use the optimizer 33 | in conjunction with: 34 | - Gradual learning rate warm-up 35 | - Linear learning rate scaling 36 | - Poly rule learning rate decay 37 | 38 | Note, LARS scaling is currently only enabled for dense tensors. Sparse tensors 39 | use the default momentum optimizer. 40 | """ 41 | 42 | def __init__( 43 | self, 44 | learning_rate, 45 | momentum=0.9, 46 | weight_decay=0.0001, 47 | # The LARS coefficient is a hyperparameter 48 | eeta=0.001, 49 | epsilon=0.0, 50 | name="LARSOptimizer", 51 | # Enable skipping variables from LARS scaling. 52 | # TODO(sameerkm): Enable a direct mechanism to pass a 53 | # subset of variables to the optimizer. 54 | skip_list=None, 55 | use_nesterov=False): 56 | """Construct a new LARS Optimizer. 57 | 58 | Args: 59 | learning_rate: A `Tensor` or floating point value. The base learning rate. 60 | momentum: A floating point value. Momentum hyperparameter. 61 | weight_decay: A floating point value. Weight decay hyperparameter. 62 | eeta: LARS coefficient as used in the paper. Dfault set to LARS 63 | coefficient from the paper. (eeta / weight_decay) determines the highest 64 | scaling factor in LARS. 65 | epsilon: Optional epsilon parameter to be set in models that have very 66 | small gradients. Default set to 0.0. 67 | name: Optional name prefix for variables and ops created by LARSOptimizer. 68 | skip_list: List of strings to enable skipping variables from LARS scaling. 69 | If any of the strings in skip_list is a subset of var.name, variable 70 | 'var' is skipped from LARS scaling. For a typical classification model 71 | with batch normalization, the skip_list is ['batch_normalization', 72 | 'bias'] 73 | use_nesterov: when set to True, nesterov momentum will be enabled 74 | 75 | Raises: 76 | ValueError: If a hyperparameter is set to a non-sensical value. 77 | """ 78 | if momentum < 0.0: 79 | raise ValueError("momentum should be positive: %s" % momentum) 80 | if weight_decay < 0.0: 81 | raise ValueError("weight_decay should be positive: %s" % weight_decay) 82 | super(LARSOptimizer, self).__init__(use_locking=False, name=name) 83 | 84 | self._learning_rate = learning_rate 85 | self._momentum = momentum 86 | self._weight_decay = weight_decay 87 | self._eeta = eeta 88 | self._epsilon = epsilon 89 | self._name = name 90 | self._skip_list = skip_list 91 | self._use_nesterov = use_nesterov 92 | 93 | def _create_slots(self, var_list): 94 | for v in var_list: 95 | self._zeros_slot(v, "momentum", self._name) 96 | 97 | def compute_lr(self, grad, var): 98 | scaled_lr = self._learning_rate 99 | if self._skip_list is None or not any(v in var.name 100 | for v in self._skip_list): 101 | w_norm = tf.norm(var, ord=2) 102 | g_norm = tf.norm(grad, ord=2) 103 | trust_ratio = tf.where( 104 | tf.math.greater(w_norm, 0), 105 | tf.where( 106 | tf.math.greater(g_norm, 0), 107 | (self._eeta * w_norm / 108 | (g_norm + self._weight_decay * w_norm + self._epsilon)), 1.0), 109 | 1.0) 110 | scaled_lr = self._learning_rate * trust_ratio 111 | # Add the weight regularization gradient 112 | grad = grad + self._weight_decay * var 113 | return scaled_lr, grad 114 | 115 | def _apply_dense(self, grad, var): 116 | scaled_lr, grad = self.compute_lr(grad, var) 117 | mom = self.get_slot(var, "momentum") 118 | return tf.raw_ops.ApplyMomentum( 119 | var, 120 | mom, 121 | tf.cast(1.0, var.dtype.base_dtype), 122 | grad * scaled_lr, 123 | self._momentum, 124 | use_locking=False, 125 | use_nesterov=self._use_nesterov) 126 | 127 | def _resource_apply_dense(self, grad, var): 128 | scaled_lr, grad = self.compute_lr(grad, var) 129 | mom = self.get_slot(var, "momentum") 130 | return tf.raw_ops.ResourceApplyMomentum( 131 | var=var.handle, 132 | accum=mom.handle, 133 | lr=tf.cast(1.0, var.dtype.base_dtype), 134 | grad=grad * scaled_lr, 135 | momentum=self._momentum, 136 | use_locking=False, 137 | use_nesterov=self._use_nesterov) 138 | 139 | # Fallback to momentum optimizer for sparse tensors 140 | def _apply_sparse(self, grad, var): 141 | mom = self.get_slot(var, "momentum") 142 | return tf.raw_ops.SparseApplyMomentum( 143 | var, 144 | mom, 145 | tf.cast(self._learning_rate_tensor, var.dtype.base_dtype), 146 | grad.values, 147 | grad.indices, 148 | tf.cast(self._momentum_tensor, var.dtype.base_dtype), 149 | use_locking=self._use_locking, 150 | use_nesterov=self._use_nesterov).op 151 | 152 | def _resource_apply_sparse(self, grad, var, indices): 153 | mom = self.get_slot(var, "momentum") 154 | return tf.raw_ops.ResourceSparseApplyMomentum( 155 | var.handle, 156 | mom.handle, 157 | tf.cast(self._learning_rate_tensor, grad.dtype), 158 | grad, 159 | indices, 160 | tf.cast(self._momentum_tensor, grad.dtype), 161 | use_locking=self._use_locking, 162 | use_nesterov=self._use_nesterov) 163 | 164 | def _prepare(self): 165 | learning_rate = self._learning_rate 166 | if callable(learning_rate): 167 | learning_rate = learning_rate() 168 | self._learning_rate_tensor = tf.convert_to_tensor( 169 | learning_rate, name="learning_rate") 170 | momentum = self._momentum 171 | if callable(momentum): 172 | momentum = momentum() 173 | self._momentum_tensor = tf.convert_to_tensor(momentum, name="momentum") 174 | -------------------------------------------------------------------------------- /efficientnet_tf/edgetpu/efficientnet_edgetpu_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Builder for EfficientNet-EdgeTPU models.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import os 22 | from absl import logging 23 | import tensorflow.compat.v1 as tf 24 | 25 | from .. import efficientnet_builder 26 | from .. import efficientnet_model 27 | from .. import utils 28 | 29 | # The input tensor is in the range of [0, 255], we need to scale them to the 30 | # range of [0, 1] 31 | MEAN_RGB = [127.0, 127.0, 127.0] 32 | STDDEV_RGB = [128.0, 128.0, 128.0] 33 | 34 | 35 | def efficientnet_edgetpu_params(model_name): 36 | """Get efficientnet-edgetpu params based on model name.""" 37 | params_dict = { 38 | # (width_coefficient, depth_coefficient, resolution, dropout_rate) 39 | 'efficientnet-edgetpu-S': (1.0, 1.0, 224, 0.2), 40 | 'efficientnet-edgetpu-M': (1.0, 1.1, 240, 0.2), 41 | 'efficientnet-edgetpu-L': (1.2, 1.4, 300, 0.3), 42 | } 43 | return params_dict[model_name] 44 | 45 | 46 | def efficientnet_edgetpu(width_coefficient=None, 47 | depth_coefficient=None, 48 | dropout_rate=0.2, 49 | survival_prob=0.8): 50 | """Creates an efficientnet-edgetpu model.""" 51 | blocks_args = [ 52 | 'r1_k3_s11_e4_i24_o24_c1_noskip', 53 | 'r2_k3_s22_e8_i24_o32_c1', 54 | 'r4_k3_s22_e8_i32_o48_c1', 55 | 'r5_k5_s22_e8_i48_o96', 56 | 'r4_k5_s11_e8_i96_o144', 57 | 'r2_k5_s22_e8_i144_o192', 58 | ] 59 | global_params = efficientnet_model.GlobalParams( 60 | batch_norm_momentum=0.99, 61 | batch_norm_epsilon=1e-3, 62 | dropout_rate=dropout_rate, 63 | survival_prob=survival_prob, 64 | data_format='channels_last', 65 | num_classes=1001, 66 | width_coefficient=width_coefficient, 67 | depth_coefficient=depth_coefficient, 68 | depth_divisor=8, 69 | min_depth=None, 70 | relu_fn=tf.nn.relu, 71 | # The default is TPU-specific batch norm. 72 | # The alternative is tf.layers.BatchNormalization. 73 | batch_norm=utils.TpuBatchNormalization, # TPU-specific requirement. 74 | local_pooling=True, # for EdgeTPU. 75 | use_se=False) 76 | decoder = efficientnet_builder.BlockDecoder() 77 | return decoder.decode(blocks_args), global_params 78 | 79 | 80 | def get_model_params(model_name, override_params): 81 | """Get the block args and global params for a given model.""" 82 | if model_name.startswith('efficientnet-edgetpu'): 83 | width_coefficient, depth_coefficient, _, dropout_rate = ( 84 | efficientnet_edgetpu_params(model_name)) 85 | blocks_args, global_params = efficientnet_edgetpu(width_coefficient, 86 | depth_coefficient, 87 | dropout_rate) 88 | else: 89 | raise NotImplementedError('model name is not pre-defined: %s' % model_name) 90 | 91 | if override_params: 92 | # ValueError will be raised here if override_params has fields not included 93 | # in global_params. 94 | global_params = global_params._replace(**override_params) 95 | 96 | logging.info('global_params= %s', global_params) 97 | logging.info('blocks_args= %s', blocks_args) 98 | return blocks_args, global_params 99 | 100 | 101 | def build_model(images, 102 | model_name, 103 | training, 104 | override_params=None, 105 | model_dir=None, 106 | fine_tuning=False): 107 | """A helper functiion to creates a model and returns predicted logits. 108 | 109 | Args: 110 | images: input images tensor. 111 | model_name: string, the predefined model name. 112 | training: boolean, whether the model is constructed for training. 113 | override_params: A dictionary of params for overriding. Fields must exist in 114 | efficientnet_model.GlobalParams. 115 | model_dir: string, optional model dir for saving configs. 116 | fine_tuning: boolean, whether the model is used for finetuning. 117 | 118 | Returns: 119 | logits: the logits tensor of classes. 120 | endpoints: the endpoints for each layer. 121 | 122 | Raises: 123 | When model_name specified an undefined model, raises NotImplementedError. 124 | When override_params has invalid fields, raises ValueError. 125 | """ 126 | assert isinstance(images, tf.Tensor) 127 | if not training or fine_tuning: 128 | if not override_params: 129 | override_params = {} 130 | override_params['batch_norm'] = utils.BatchNormalization 131 | blocks_args, global_params = get_model_params(model_name, override_params) 132 | if not training or fine_tuning: 133 | global_params = global_params._replace(batch_norm=utils.BatchNormalization) 134 | 135 | if model_dir: 136 | param_file = os.path.join(model_dir, 'model_params.txt') 137 | if not tf.gfile.Exists(param_file): 138 | if not tf.gfile.Exists(model_dir): 139 | tf.gfile.MakeDirs(model_dir) 140 | with tf.gfile.GFile(param_file, 'w') as f: 141 | logging.info('writing to %s', param_file) 142 | f.write('model_name= %s\n\n' % model_name) 143 | f.write('global_params= %s\n\n' % str(global_params)) 144 | f.write('blocks_args= %s\n\n' % str(blocks_args)) 145 | 146 | with tf.variable_scope(model_name): 147 | model = efficientnet_model.Model(blocks_args, global_params) 148 | logits = model(images, training=training) 149 | 150 | logits = tf.identity(logits, 'logits') 151 | return logits, model.endpoints 152 | 153 | 154 | def build_model_base(images, model_name, training, override_params=None): 155 | """A helper functiion to create a base model and return global_pool. 156 | 157 | Args: 158 | images: input images tensor. 159 | model_name: string, the model name of a pre-defined MnasNet. 160 | training: boolean, whether the model is constructed for training. 161 | override_params: A dictionary of params for overriding. Fields must exist in 162 | mnasnet_model.GlobalParams. 163 | 164 | Returns: 165 | features: global pool features. 166 | endpoints: the endpoints for each layer. 167 | 168 | Raises: 169 | When model_name specified an undefined model, raises NotImplementedError. 170 | When override_params has invalid fields, raises ValueError. 171 | """ 172 | assert isinstance(images, tf.Tensor) 173 | blocks_args, global_params = get_model_params(model_name, override_params) 174 | 175 | with tf.variable_scope(model_name): 176 | model = efficientnet_model.Model(blocks_args, global_params) 177 | features = model(images, training=training, features_only=True) 178 | 179 | features = tf.identity(features, 'global_pool') 180 | return features, model.endpoints 181 | -------------------------------------------------------------------------------- /efficientnet_tf/condconv/efficientnet_condconv_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Builder for EfficientNet-CondConv models.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import os 22 | import tensorflow.compat.v1 as tf 23 | 24 | from .. import efficientnet_builder 25 | from .. import efficientnet_model 26 | from .. import utils 27 | 28 | # The input tensor is in the range of [0, 255], we need to scale them to the 29 | # range of [0, 1] 30 | MEAN_RGB = [127.0, 127.0, 127.0] 31 | STDDEV_RGB = [128.0, 128.0, 128.0] 32 | 33 | 34 | def efficientnet_condconv_params(model_name): 35 | """Get efficientnet-condconv params based on model name.""" 36 | params_dict = { 37 | # (width_coefficient, depth_coefficient, resolution, dropout_rate, 38 | # condconv_num_experts) 39 | 'efficientnet-condconv-b0-4e': (1.0, 1.0, 224, 0.25, 4), 40 | 'efficientnet-condconv-b0-8e': (1.0, 1.0, 224, 0.25, 8), 41 | 'efficientnet-condconv-b0-8e-depth': (1.0, 1.1, 224, 0.25, 8) 42 | } 43 | return params_dict[model_name] 44 | 45 | 46 | def efficientnet_condconv(width_coefficient=None, 47 | depth_coefficient=None, 48 | dropout_rate=0.2, 49 | survival_prob=0.8, 50 | condconv_num_experts=None): 51 | """Creates an efficientnet-condconv model.""" 52 | blocks_args = [ 53 | 'r1_k3_s11_e1_i32_o16_se0.25', 54 | 'r2_k3_s22_e6_i16_o24_se0.25', 55 | 'r2_k5_s22_e6_i24_o40_se0.25', 56 | 'r3_k3_s22_e6_i40_o80_se0.25', 57 | 'r3_k5_s11_e6_i80_o112_se0.25_cc', 58 | 'r4_k5_s22_e6_i112_o192_se0.25_cc', 59 | 'r1_k3_s11_e6_i192_o320_se0.25_cc', 60 | ] 61 | global_params = efficientnet_model.GlobalParams( 62 | batch_norm_momentum=0.99, 63 | batch_norm_epsilon=1e-3, 64 | dropout_rate=dropout_rate, 65 | survival_prob=survival_prob, 66 | data_format='channels_last', 67 | num_classes=1000, 68 | width_coefficient=width_coefficient, 69 | depth_coefficient=depth_coefficient, 70 | depth_divisor=8, 71 | min_depth=None, 72 | relu_fn=tf.nn.swish, 73 | # The default is TPU-specific batch norm. 74 | # The alternative is tf.layers.BatchNormalization. 75 | batch_norm=utils.TpuBatchNormalization, # TPU-specific requirement. 76 | use_se=True, 77 | condconv_num_experts=condconv_num_experts) 78 | decoder = efficientnet_builder.BlockDecoder() 79 | return decoder.decode(blocks_args), global_params 80 | 81 | 82 | def get_model_params(model_name, override_params): 83 | """Get the block args and global params for a given model.""" 84 | if model_name.startswith('efficientnet-condconv'): 85 | (width_coefficient, depth_coefficient, _, dropout_rate, 86 | condconv_num_experts) = ( 87 | efficientnet_condconv_params(model_name)) 88 | blocks_args, global_params = efficientnet_condconv( 89 | width_coefficient=width_coefficient, 90 | depth_coefficient=depth_coefficient, 91 | dropout_rate=dropout_rate, 92 | condconv_num_experts=condconv_num_experts) 93 | else: 94 | raise NotImplementedError('model name is not pre-defined: %s' % model_name) 95 | 96 | if override_params: 97 | # ValueError will be raised here if override_params has fields not included 98 | # in global_params. 99 | global_params = global_params._replace(**override_params) 100 | 101 | tf.logging.info('global_params= %s', global_params) 102 | tf.logging.info('blocks_args= %s', blocks_args) 103 | return blocks_args, global_params 104 | 105 | 106 | def build_model(images, 107 | model_name, 108 | training, 109 | override_params=None, 110 | model_dir=None, 111 | fine_tuning=False): 112 | """A helper functiion to creates a model and returns predicted logits. 113 | 114 | Args: 115 | images: input images tensor. 116 | model_name: string, the predefined model name. 117 | training: boolean, whether the model is constructed for training. 118 | override_params: A dictionary of params for overriding. Fields must exist in 119 | efficientnet_model.GlobalParams. 120 | model_dir: string, optional model dir for saving configs. 121 | fine_tuning: boolean, whether the model is used for finetuning. 122 | 123 | Returns: 124 | logits: the logits tensor of classes. 125 | endpoints: the endpoints for each layer. 126 | 127 | Raises: 128 | When model_name specified an undefined model, raises NotImplementedError. 129 | When override_params has invalid fields, raises ValueError. 130 | """ 131 | assert isinstance(images, tf.Tensor) 132 | if not training or fine_tuning: 133 | if not override_params: 134 | override_params = {} 135 | override_params['batch_norm'] = utils.BatchNormalization 136 | blocks_args, global_params = get_model_params(model_name, override_params) 137 | if not training or fine_tuning: 138 | global_params = global_params._replace(batch_norm=utils.BatchNormalization) 139 | 140 | if model_dir: 141 | param_file = os.path.join(model_dir, 'model_params.txt') 142 | if not tf.gfile.Exists(param_file): 143 | if not tf.gfile.Exists(model_dir): 144 | tf.gfile.MakeDirs(model_dir) 145 | with tf.gfile.GFile(param_file, 'w') as f: 146 | tf.logging.info('writing to %s' % param_file) 147 | f.write('model_name= %s\n\n' % model_name) 148 | f.write('global_params= %s\n\n' % str(global_params)) 149 | f.write('blocks_args= %s\n\n' % str(blocks_args)) 150 | 151 | with tf.variable_scope(model_name): 152 | model = efficientnet_model.Model(blocks_args, global_params) 153 | logits = model(images, training=training) 154 | 155 | logits = tf.identity(logits, 'logits') 156 | return logits, model.endpoints 157 | 158 | 159 | def build_model_base(images, model_name, training, override_params=None): 160 | """A helper functiion to create a base model and return global_pool. 161 | 162 | Args: 163 | images: input images tensor. 164 | model_name: string, the model name of a pre-defined MnasNet. 165 | training: boolean, whether the model is constructed for training. 166 | override_params: A dictionary of params for overriding. Fields must exist in 167 | mnasnet_model.GlobalParams. 168 | 169 | Returns: 170 | features: global pool features. 171 | endpoints: the endpoints for each layer. 172 | 173 | Raises: 174 | When model_name specified an undefined model, raises NotImplementedError. 175 | When override_params has invalid fields, raises ValueError. 176 | """ 177 | assert isinstance(images, tf.Tensor) 178 | blocks_args, global_params = get_model_params(model_name, override_params) 179 | 180 | with tf.variable_scope(model_name): 181 | model = efficientnet_model.Model(blocks_args, global_params) 182 | features = model(images, training=training, features_only=True) 183 | 184 | features = tf.identity(features, 'global_pool') 185 | return features, model.endpoints 186 | -------------------------------------------------------------------------------- /efficientnet_tf/tpu/efficientnet_x_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Model Builder for EfficientNet-X.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import functools 22 | import os 23 | 24 | from absl import logging 25 | import tensorflow.compat.v1 as tf 26 | 27 | from .. import efficientnet_builder 28 | from .. import efficientnet_model 29 | from .. import utils 30 | 31 | MEAN_RGB = [0.485 * 255, 0.456 * 255, 0.406 * 255] 32 | STDDEV_RGB = [0.229 * 255, 0.224 * 255, 0.225 * 255] 33 | 34 | 35 | def efficientnet_x_params(model_name): 36 | """Get efficientnet params based on model name.""" 37 | params_dict = { 38 | # (width_coefficient, depth_coefficient, resolution, dropout_rate, 39 | # se_coefficient) 40 | 'efficientnet-x-b0': (1.0, 1.0, 224, 0.2, 4), 41 | 'efficientnet-x-b1': (1.0, 1.1, 240, 0.2, 2), 42 | 'efficientnet-x-b2': (1.1, 1.2, 260, 0.3, 1), 43 | 'efficientnet-x-b3': (1.2, 1.4, 300, 0.3, 1), 44 | 'efficientnet-x-b4': (1.4, 1.8, 380, 0.4, 1), 45 | 'efficientnet-x-b5': (1.6, 2.2, 456, 0.4, 1), 46 | 'efficientnet-x-b6': (1.8, 2.6, 528, 0.5, 1), 47 | 'efficientnet-x-b7': (2.0, 3.1, 600, 0.5, 1), 48 | } 49 | return params_dict[model_name] 50 | 51 | 52 | def efficientnet_x(width_coefficient=None, 53 | depth_coefficient=None, 54 | se_coefficient=None, 55 | dropout_rate=0.2, 56 | survival_prob=0.8): 57 | """Creates a efficientnet model.""" 58 | blocks_args = [ 59 | 'r1_k3_s11_e1_i32_o16_se0.25_d1_a0', 60 | 'r2_k3_s22_e6_i16_o24_se0.25_f1_d2_a1', 61 | 'r2_k5_s22_e6_i24_o40_se0.25_f1_a1', 62 | 'r3_k3_s22_e6_i40_o80_se0.25_a0', 63 | 'r3_k5_s11_e6_i80_o112_se0.25_a0', 64 | 'r4_k5_s22_e6_i112_o192_se0.25_a0', 65 | 'r1_k3_s11_e6_i192_o320_se0.25_a0', 66 | ] 67 | global_params = efficientnet_model.GlobalParams( 68 | batch_norm_momentum=0.99, 69 | batch_norm_epsilon=1e-3, 70 | dropout_rate=dropout_rate, 71 | survival_prob=survival_prob, 72 | data_format='channels_last', 73 | num_classes=1000, 74 | width_coefficient=width_coefficient, 75 | depth_coefficient=depth_coefficient, 76 | depth_divisor=8, 77 | min_depth=None, 78 | relu_fn=tf.nn.relu, 79 | # The default is TPU-specific batch norm. 80 | # The alternative is tf.layers.BatchNormalization. 81 | batch_norm=utils.TpuBatchNormalization, # TPU-specific requirement. 82 | use_se=True, 83 | se_coefficient=se_coefficient) 84 | decoder = efficientnet_builder.BlockDecoder() 85 | return decoder.decode(blocks_args), global_params 86 | 87 | 88 | def get_model_params(model_name, override_params): 89 | """Get the block args and global params for a given model.""" 90 | if model_name.startswith('efficientnet'): 91 | width_coefficient, depth_coefficient, _, dropout_rate, se_coefficient = ( 92 | efficientnet_x_params(model_name)) 93 | blocks_args, global_params = efficientnet_x( 94 | width_coefficient, depth_coefficient, se_coefficient, dropout_rate) 95 | else: 96 | raise NotImplementedError('model name is not pre-defined: %s' % model_name) 97 | 98 | if override_params: 99 | # ValueError will be raised here if override_params has fields not included 100 | # in global_params. 101 | global_params = global_params._replace(**override_params) 102 | 103 | logging.info('global_params= %s', global_params) 104 | logging.info('blocks_args= %s', blocks_args) 105 | return blocks_args, global_params 106 | 107 | 108 | def build_model(images, 109 | model_name, 110 | training, 111 | override_params=None, 112 | model_dir=None, 113 | fine_tuning=False, 114 | features_only=False, 115 | pooled_features_only=False): 116 | """A helper function to creates a model and returns predicted logits. 117 | 118 | Args: 119 | images: input images tensor. 120 | model_name: string, the predefined model name. 121 | training: boolean, whether the model is constructed for training. 122 | override_params: A dictionary of params for overriding. Fields must exist in 123 | efficientnet_model.GlobalParams. 124 | model_dir: string, optional model dir for saving configs. 125 | fine_tuning: boolean, whether the model is used for finetuning. 126 | features_only: build the base feature network only (excluding final 127 | 1x1 conv layer, global pooling, dropout and fc head). 128 | pooled_features_only: build the base network for features extraction (after 129 | 1x1 conv layer and global pooling, but before dropout and fc head). 130 | 131 | Returns: 132 | logits: the logits tensor of classes. 133 | endpoints: the endpoints for each layer. 134 | 135 | Raises: 136 | When model_name specified an undefined model, raises NotImplementedError. 137 | When override_params has invalid fields, raises ValueError. 138 | """ 139 | assert isinstance(images, tf.Tensor) 140 | assert not (features_only and pooled_features_only) 141 | if not training or fine_tuning: 142 | if not override_params: 143 | override_params = {} 144 | override_params['batch_norm'] = utils.BatchNormalization 145 | if fine_tuning: 146 | override_params['relu_fn'] = functools.partial( 147 | efficientnet_builder.swish, use_native=False) 148 | blocks_args, global_params = get_model_params(model_name, override_params) 149 | 150 | if model_dir: 151 | param_file = os.path.join(model_dir, 'model_params.txt') 152 | if not tf.gfile.Exists(param_file): 153 | if not tf.gfile.Exists(model_dir): 154 | tf.gfile.MakeDirs(model_dir) 155 | with tf.gfile.GFile(param_file, 'w') as f: 156 | logging.info('writing to %s', param_file) 157 | f.write('model_name= %s\n\n' % model_name) 158 | f.write('global_params= %s\n\n' % str(global_params)) 159 | f.write('blocks_args= %s\n\n' % str(blocks_args)) 160 | 161 | with tf.variable_scope(model_name): 162 | model = efficientnet_model.Model(blocks_args, global_params) 163 | outputs = model( 164 | images, 165 | training=training, 166 | features_only=features_only, 167 | pooled_features_only=pooled_features_only) 168 | if features_only: 169 | outputs = tf.identity(outputs, 'features') 170 | elif pooled_features_only: 171 | outputs = tf.identity(outputs, 'pooled_features') 172 | else: 173 | outputs = tf.identity(outputs, 'logits') 174 | return outputs, model.endpoints 175 | 176 | 177 | def build_model_base(images, model_name, training, override_params=None): 178 | """Create a base feature network and return the features before pooling. 179 | 180 | Args: 181 | images: input images tensor. 182 | model_name: string, the predefined model name. 183 | training: boolean, whether the model is constructed for training. 184 | override_params: A dictionary of params for overriding. Fields must exist in 185 | efficientnet_model.GlobalParams. 186 | 187 | Returns: 188 | features: base features before pooling. 189 | endpoints: the endpoints for each layer. 190 | 191 | Raises: 192 | When model_name specified an undefined model, raises NotImplementedError. 193 | When override_params has invalid fields, raises ValueError. 194 | """ 195 | assert isinstance(images, tf.Tensor) 196 | blocks_args, global_params = get_model_params(model_name, override_params) 197 | 198 | with tf.variable_scope(model_name): 199 | model = efficientnet_model.Model(blocks_args, global_params) 200 | features = model(images, training=training, features_only=True) 201 | 202 | features = tf.identity(features, 'features') 203 | return features, model.endpoints 204 | -------------------------------------------------------------------------------- /efficientnet_tf/lite/efficientnet_lite_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Model Builder for EfficientNet Edge Models.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import os 22 | from absl import logging 23 | import tensorflow.compat.v1 as tf 24 | 25 | from .. import efficientnet_builder 26 | from .. import efficientnet_model 27 | from .. import utils 28 | # Edge models use inception-style MEAN and STDDEV for better post-quantization. 29 | MEAN_RGB = [127.0, 127.0, 127.0] 30 | STDDEV_RGB = [128.0, 128.0, 128.0] 31 | 32 | 33 | def efficientnet_lite_params(model_name): 34 | """Get efficientnet params based on model name.""" 35 | params_dict = { 36 | # (width_coefficient, depth_coefficient, resolution, dropout_rate) 37 | 'efficientnet-lite0': (1.0, 1.0, 224, 0.2), 38 | 'efficientnet-lite1': (1.0, 1.1, 240, 0.2), 39 | 'efficientnet-lite2': (1.1, 1.2, 260, 0.3), 40 | 'efficientnet-lite3': (1.2, 1.4, 280, 0.3), 41 | 'efficientnet-lite4': (1.4, 1.8, 300, 0.3), 42 | } 43 | return params_dict[model_name] 44 | 45 | 46 | _DEFAULT_BLOCKS_ARGS = [ 47 | 'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25', 48 | 'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25', 49 | 'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25', 50 | 'r1_k3_s11_e6_i192_o320_se0.25', 51 | ] 52 | 53 | 54 | def efficientnet_lite(width_coefficient=None, 55 | depth_coefficient=None, 56 | dropout_rate=0.2, 57 | survival_prob=0.8): 58 | """Creates a efficientnet model.""" 59 | global_params = efficientnet_model.GlobalParams( 60 | blocks_args=_DEFAULT_BLOCKS_ARGS, 61 | batch_norm_momentum=0.99, 62 | batch_norm_epsilon=1e-3, 63 | dropout_rate=dropout_rate, 64 | survival_prob=survival_prob, 65 | data_format='channels_last', 66 | num_classes=1000, 67 | width_coefficient=width_coefficient, 68 | depth_coefficient=depth_coefficient, 69 | depth_divisor=8, 70 | min_depth=None, 71 | relu_fn=tf.nn.relu6, # Relu6 is for easier quantization. 72 | # The default is TPU-specific batch norm. 73 | # The alternative is tf.layers.BatchNormalization. 74 | batch_norm=utils.TpuBatchNormalization, # TPU-specific requirement. 75 | clip_projection_output=False, 76 | fix_head_stem=True, # Don't scale stem and head. 77 | local_pooling=True, # special cases for tflite issues. 78 | use_se=False) # SE is not well supported on many lite devices. 79 | return global_params 80 | 81 | 82 | def get_model_params(model_name, override_params): 83 | """Get the block args and global params for a given model.""" 84 | if model_name.startswith('efficientnet-lite'): 85 | width_coefficient, depth_coefficient, _, dropout_rate = ( 86 | efficientnet_lite_params(model_name)) 87 | global_params = efficientnet_lite( 88 | width_coefficient, depth_coefficient, dropout_rate) 89 | else: 90 | raise NotImplementedError('model name is not pre-defined: %s' % model_name) 91 | 92 | if override_params: 93 | # ValueError will be raised here if override_params has fields not included 94 | # in global_params. 95 | global_params = global_params._replace(**override_params) 96 | 97 | decoder = efficientnet_builder.BlockDecoder() 98 | blocks_args = decoder.decode(global_params.blocks_args) 99 | 100 | logging.info('global_params= %s', global_params) 101 | return blocks_args, global_params 102 | 103 | 104 | def build_model(images, 105 | model_name, 106 | training, 107 | override_params=None, 108 | model_dir=None, 109 | fine_tuning=False, 110 | features_only=False, 111 | pooled_features_only=False): 112 | """A helper function to create a model and return predicted logits. 113 | 114 | Args: 115 | images: input images tensor. 116 | model_name: string, the predefined model name. 117 | training: boolean, whether the model is constructed for training. 118 | override_params: A dictionary of params for overriding. Fields must exist in 119 | efficientnet_model.GlobalParams. 120 | model_dir: string, optional model dir for saving configs. 121 | fine_tuning: boolean, whether the model is used for finetuning. 122 | features_only: build the base feature network only (excluding final 123 | 1x1 conv layer, global pooling, dropout and fc head). 124 | pooled_features_only: build the base network for features extraction (after 125 | 1x1 conv layer and global pooling, but before dropout and fc head). 126 | 127 | Returns: 128 | logits: the logits tensor of classes. 129 | endpoints: the endpoints for each layer. 130 | 131 | Raises: 132 | When model_name specified an undefined model, raises NotImplementedError. 133 | When override_params has invalid fields, raises ValueError. 134 | """ 135 | assert isinstance(images, tf.Tensor) 136 | assert not (features_only and pooled_features_only) 137 | 138 | # For backward compatibility. 139 | if override_params and override_params.get('drop_connect_rate', None): 140 | override_params['survival_prob'] = 1 - override_params['drop_connect_rate'] 141 | 142 | if not training or fine_tuning: 143 | if not override_params: 144 | override_params = {} 145 | override_params['batch_norm'] = utils.BatchNormalization 146 | blocks_args, global_params = get_model_params(model_name, override_params) 147 | 148 | if model_dir: 149 | param_file = os.path.join(model_dir, 'model_params.txt') 150 | if not tf.gfile.Exists(param_file): 151 | if not tf.gfile.Exists(model_dir): 152 | tf.gfile.MakeDirs(model_dir) 153 | with tf.gfile.GFile(param_file, 'w') as f: 154 | logging.info('writing to %s', param_file) 155 | f.write('model_name= %s\n\n' % model_name) 156 | f.write('global_params= %s\n\n' % str(global_params)) 157 | f.write('blocks_args= %s\n\n' % str(blocks_args)) 158 | 159 | with tf.variable_scope(model_name): 160 | model = efficientnet_model.Model(blocks_args, global_params) 161 | outputs = model( 162 | images, 163 | training=training, 164 | features_only=features_only, 165 | pooled_features_only=pooled_features_only) 166 | if features_only: 167 | outputs = tf.identity(outputs, 'features') 168 | elif pooled_features_only: 169 | outputs = tf.identity(outputs, 'pooled_features') 170 | else: 171 | outputs = tf.identity(outputs, 'logits') 172 | return outputs, model.endpoints 173 | 174 | 175 | def build_model_base(images, model_name, training, override_params=None): 176 | """Create a base feature network and return the features before pooling. 177 | 178 | Args: 179 | images: input images tensor. 180 | model_name: string, the predefined model name. 181 | training: boolean, whether the model is constructed for training. 182 | override_params: A dictionary of params for overriding. Fields must exist in 183 | efficientnet_model.GlobalParams. 184 | 185 | Returns: 186 | features: base features before pooling. 187 | endpoints: the endpoints for each layer. 188 | 189 | Raises: 190 | When model_name specified an undefined model, raises NotImplementedError. 191 | When override_params has invalid fields, raises ValueError. 192 | """ 193 | assert isinstance(images, tf.Tensor) 194 | # For backward compatibility. 195 | if override_params and override_params.get('drop_connect_rate', None): 196 | override_params['survival_prob'] = 1 - override_params['drop_connect_rate'] 197 | 198 | blocks_args, global_params = get_model_params(model_name, override_params) 199 | 200 | with tf.variable_scope(model_name): 201 | model = efficientnet_model.Model(blocks_args, global_params) 202 | features = model(images, training=training, features_only=True) 203 | 204 | features = tf.identity(features, 'features') 205 | return features, model.endpoints 206 | -------------------------------------------------------------------------------- /efficientnet_tf/preprocessing.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """ImageNet preprocessing.""" 16 | from __future__ import absolute_import 17 | from __future__ import division 18 | from __future__ import print_function 19 | 20 | from absl import logging 21 | 22 | import tensorflow.compat.v1 as tf 23 | 24 | 25 | IMAGE_SIZE = 224 26 | CROP_PADDING = 32 27 | 28 | 29 | def distorted_bounding_box_crop(image_bytes, 30 | bbox, 31 | min_object_covered=0.1, 32 | aspect_ratio_range=(0.75, 1.33), 33 | area_range=(0.05, 1.0), 34 | max_attempts=100, 35 | scope=None): 36 | """Generates cropped_image using one of the bboxes randomly distorted. 37 | 38 | See `tf.image.sample_distorted_bounding_box` for more documentation. 39 | 40 | Args: 41 | image_bytes: `Tensor` of binary image data. 42 | bbox: `Tensor` of bounding boxes arranged `[1, num_boxes, coords]` 43 | where each coordinate is [0, 1) and the coordinates are arranged 44 | as `[ymin, xmin, ymax, xmax]`. If num_boxes is 0 then use the whole 45 | image. 46 | min_object_covered: An optional `float`. Defaults to `0.1`. The cropped 47 | area of the image must contain at least this fraction of any bounding 48 | box supplied. 49 | aspect_ratio_range: An optional list of `float`s. The cropped area of the 50 | image must have an aspect ratio = width / height within this range. 51 | area_range: An optional list of `float`s. The cropped area of the image 52 | must contain a fraction of the supplied image within in this range. 53 | max_attempts: An optional `int`. Number of attempts at generating a cropped 54 | region of the image of the specified constraints. After `max_attempts` 55 | failures, return the entire image. 56 | scope: Optional `str` for name scope. 57 | Returns: 58 | cropped image `Tensor` 59 | """ 60 | with tf.name_scope(scope, 'distorted_bounding_box_crop', [image_bytes, bbox]): 61 | shape = tf.image.extract_jpeg_shape(image_bytes) 62 | sample_distorted_bounding_box = tf.image.sample_distorted_bounding_box( 63 | shape, 64 | bounding_boxes=bbox, 65 | min_object_covered=min_object_covered, 66 | aspect_ratio_range=aspect_ratio_range, 67 | area_range=area_range, 68 | max_attempts=max_attempts, 69 | use_image_if_no_bounding_boxes=True) 70 | bbox_begin, bbox_size, _ = sample_distorted_bounding_box 71 | 72 | # Crop the image to the specified bounding box. 73 | offset_y, offset_x, _ = tf.unstack(bbox_begin) 74 | target_height, target_width, _ = tf.unstack(bbox_size) 75 | crop_window = tf.stack([offset_y, offset_x, target_height, target_width]) 76 | image = tf.image.decode_and_crop_jpeg(image_bytes, crop_window, channels=3) 77 | 78 | return image 79 | 80 | 81 | def _at_least_x_are_equal(a, b, x): 82 | """At least `x` of `a` and `b` `Tensors` are equal.""" 83 | match = tf.equal(a, b) 84 | match = tf.cast(match, tf.int32) 85 | return tf.greater_equal(tf.reduce_sum(match), x) 86 | 87 | 88 | def _resize_image(image, image_size, method=None): 89 | if method is not None: 90 | tf.logging.info('Use customized resize method {}'.format(method)) 91 | return tf.image.resize([image], [image_size, image_size], method)[0] 92 | tf.logging.info('Use default resize_bicubic.') 93 | return tf.image.resize_bicubic([image], [image_size, image_size])[0] 94 | 95 | 96 | def _decode_and_random_crop(image_bytes, image_size, resize_method=None): 97 | """Make a random crop of image_size.""" 98 | bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4]) 99 | image = distorted_bounding_box_crop( 100 | image_bytes, 101 | bbox, 102 | min_object_covered=0.1, 103 | aspect_ratio_range=(3. / 4, 4. / 3.), 104 | area_range=(0.08, 1.0), 105 | max_attempts=10, 106 | scope=None) 107 | original_shape = tf.image.extract_jpeg_shape(image_bytes) 108 | bad = _at_least_x_are_equal(original_shape, tf.shape(image), 3) 109 | 110 | image = tf.cond( 111 | bad, 112 | lambda: _decode_and_center_crop(image_bytes, image_size), 113 | lambda: _resize_image(image, image_size, resize_method)) 114 | 115 | return image 116 | 117 | 118 | def _decode_and_center_crop(image_bytes, image_size, resize_method=None): 119 | """Crops to center of image with padding then scales image_size.""" 120 | shape = tf.image.extract_jpeg_shape(image_bytes) 121 | image_height = shape[0] 122 | image_width = shape[1] 123 | 124 | padded_center_crop_size = tf.cast( 125 | ((image_size / (image_size + CROP_PADDING)) * 126 | tf.cast(tf.minimum(image_height, image_width), tf.float32)), 127 | tf.int32) 128 | 129 | offset_height = ((image_height - padded_center_crop_size) + 1) // 2 130 | offset_width = ((image_width - padded_center_crop_size) + 1) // 2 131 | crop_window = tf.stack([offset_height, offset_width, 132 | padded_center_crop_size, padded_center_crop_size]) 133 | image = tf.image.decode_and_crop_jpeg(image_bytes, crop_window, channels=3) 134 | image = _resize_image(image, image_size, resize_method) 135 | return image 136 | 137 | 138 | def _flip(image): 139 | """Random horizontal image flip.""" 140 | image = tf.image.random_flip_left_right(image) 141 | return image 142 | 143 | 144 | def preprocess_for_train(image_bytes, 145 | use_bfloat16, 146 | image_size=IMAGE_SIZE, 147 | augment_name=None, 148 | randaug_num_layers=None, 149 | randaug_magnitude=None, 150 | resize_method=None): 151 | """Preprocesses the given image for evaluation. 152 | 153 | Args: 154 | image_bytes: `Tensor` representing an image binary of arbitrary size. 155 | use_bfloat16: `bool` for whether to use bfloat16. 156 | image_size: image size. 157 | augment_name: `string` that is the name of the augmentation method 158 | to apply to the image. `autoaugment` if AutoAugment is to be used or 159 | `randaugment` if RandAugment is to be used. If the value is `None` no 160 | augmentation method will be applied applied. See autoaugment.py for more 161 | details. 162 | randaug_num_layers: 'int', if RandAug is used, what should the number of 163 | layers be. See autoaugment.py for detailed description. 164 | randaug_magnitude: 'int', if RandAug is used, what should the magnitude 165 | be. See autoaugment.py for detailed description. 166 | resize_method: resize method. If none, use bicubic. 167 | 168 | Returns: 169 | A preprocessed image `Tensor`. 170 | """ 171 | image = _decode_and_random_crop(image_bytes, image_size, resize_method) 172 | image = _flip(image) 173 | image = tf.reshape(image, [image_size, image_size, 3]) 174 | 175 | if augment_name: 176 | try: 177 | import autoaugment # pylint: disable=g-import-not-at-top 178 | except ImportError as e: 179 | logging.exception('Autoaugment is not supported in TF 2.x.') 180 | raise e 181 | 182 | logging.info('Apply AutoAugment policy %s', augment_name) 183 | input_image_type = image.dtype 184 | image = tf.clip_by_value(image, 0.0, 255.0) 185 | image = tf.cast(image, dtype=tf.uint8) 186 | 187 | if augment_name == 'autoaugment': 188 | logging.info('Apply AutoAugment policy %s', augment_name) 189 | image = autoaugment.distort_image_with_autoaugment(image, 'v0') 190 | elif augment_name == 'randaugment': 191 | image = autoaugment.distort_image_with_randaugment( 192 | image, randaug_num_layers, randaug_magnitude) 193 | else: 194 | raise ValueError('Invalid value for augment_name: %s' % (augment_name)) 195 | 196 | image = tf.cast(image, dtype=input_image_type) 197 | 198 | image = tf.image.convert_image_dtype( 199 | image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32) 200 | 201 | return image 202 | 203 | 204 | def preprocess_for_eval(image_bytes, 205 | use_bfloat16, 206 | image_size=IMAGE_SIZE, 207 | resize_method=None): 208 | """Preprocesses the given image for evaluation. 209 | 210 | Args: 211 | image_bytes: `Tensor` representing an image binary of arbitrary size. 212 | use_bfloat16: `bool` for whether to use bfloat16. 213 | image_size: image size. 214 | resize_method: if None, use bicubic. 215 | 216 | Returns: 217 | A preprocessed image `Tensor`. 218 | """ 219 | image = _decode_and_center_crop(image_bytes, image_size, resize_method) 220 | image = tf.reshape(image, [image_size, image_size, 3]) 221 | image = tf.image.convert_image_dtype( 222 | image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32) 223 | return image 224 | 225 | 226 | def preprocess_image(image_bytes, 227 | is_training=False, 228 | use_bfloat16=False, 229 | image_size=IMAGE_SIZE, 230 | augment_name=None, 231 | randaug_num_layers=None, 232 | randaug_magnitude=None, 233 | resize_method=None): 234 | """Preprocesses the given image. 235 | 236 | Args: 237 | image_bytes: `Tensor` representing an image binary of arbitrary size. 238 | is_training: `bool` for whether the preprocessing is for training. 239 | use_bfloat16: `bool` for whether to use bfloat16. 240 | image_size: image size. 241 | augment_name: `string` that is the name of the augmentation method 242 | to apply to the image. `autoaugment` if AutoAugment is to be used or 243 | `randaugment` if RandAugment is to be used. If the value is `None` no 244 | augmentation method will be applied applied. See autoaugment.py for more 245 | details. 246 | randaug_num_layers: 'int', if RandAug is used, what should the number of 247 | layers be. See autoaugment.py for detailed description. 248 | randaug_magnitude: 'int', if RandAug is used, what should the magnitude 249 | be. See autoaugment.py for detailed description. 250 | resize_method: 'string' or None. Use resize_bicubic in default. 251 | 252 | Returns: 253 | A preprocessed image `Tensor` with value range of [0, 255]. 254 | """ 255 | if is_training: 256 | return preprocess_for_train( 257 | image_bytes, use_bfloat16, image_size, augment_name, 258 | randaug_num_layers, randaug_magnitude, resize_method) 259 | else: 260 | return preprocess_for_eval(image_bytes, use_bfloat16, image_size, 261 | resize_method) 262 | -------------------------------------------------------------------------------- /efficientnet_tf/README.md: -------------------------------------------------------------------------------- 1 | # EfficientNets 2 | 3 | [1] Mingxing Tan and Quoc V. Le. EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks. ICML 2019. 4 | Arxiv link: https://arxiv.org/abs/1905.11946. 5 | 6 | Updates 7 | 8 | - **[Mar 2020] Released mobile/IoT device friendly EfficientNet-lite models: [README](lite/README.md).** 9 | 10 | - [Feb 2020] Released EfficientNet checkpoints trained with NoisyStudent: [paper](https://arxiv.org/abs/1911.04252). 11 | 12 | - [Nov 2019] Released EfficientNet checkpoints trained with AdvProp: [paper](https://arxiv.org/abs/1911.09665). 13 | 14 | - [Oct 2019] Released EfficientNet-CondConv models with conditionally parameterized convolutions: [README](condconv/README.md), [paper](https://arxiv.org/abs/1904.04971). 15 | 16 | - [Oct 2019] Released EfficientNet models trained with RandAugment: [paper](https://arxiv.org/abs/1909.13719). 17 | 18 | - [Aug 2019] Released EfficientNet-EdgeTPU models: [README](edgetpu/README.md) and [blog post](https://ai.googleblog.com/2019/08/efficientnet-edgetpu-creating.html). 19 | 20 | - [Jul 2019] Released EfficientNet checkpoints trained with AutoAugment: [paper](https://arxiv.org/abs/1805.09501), [blog post](https://ai.googleblog.com/2018/06/improving-deep-learning-performance.html) 21 | 22 | - [May 2019] Released EfficientNets code and weights: [blog post](https://ai.googleblog.com/2019/05/efficientnet-improving-accuracy-and.html) 23 | 24 | ## 1. About EfficientNet Models 25 | 26 | EfficientNets are a family of image classification models, which achieve state-of-the-art accuracy, yet being an order-of-magnitude smaller and faster than previous models. 27 | 28 | We develop EfficientNets based on AutoML and Compound Scaling. In particular, we first use [AutoML MNAS Mobile framework](https://ai.googleblog.com/2018/08/mnasnet-towards-automating-design-of.html) to develop a mobile-size baseline network, named as EfficientNet-B0; Then, we use the compound scaling method to scale up this baseline to obtain EfficientNet-B1 to B7. 29 | 30 | 31 | 32 | 35 | 38 | 39 |
33 | 34 | 36 | 37 |
40 | 41 | EfficientNets achieve state-of-the-art accuracy on ImageNet with an order of magnitude better efficiency: 42 | 43 | 44 | * In high-accuracy regime, our EfficientNet-B7 achieves state-of-the-art 84.4% top-1 / 97.1% top-5 accuracy on ImageNet with 66M parameters and 37B FLOPS, being 8.4x smaller and 6.1x faster on CPU inference than previous best [Gpipe](https://arxiv.org/abs/1811.06965). 45 | 46 | * In middle-accuracy regime, our EfficientNet-B1 is 7.6x smaller and 5.7x faster on CPU inference than [ResNet-152](https://arxiv.org/abs/1512.03385), with similar ImageNet accuracy. 47 | 48 | * Compared with the widely used [ResNet-50](https://arxiv.org/abs/1512.03385), our EfficientNet-B4 improves the top-1 accuracy from 76.3% of ResNet-50 to 82.6% (+6.3%), under similar FLOPS constraint. 49 | 50 | ## 2. Using Pretrained EfficientNet Checkpoints 51 | 52 | To train EfficientNet on ImageNet, we hold out 25,022 randomly picked images ([image filenames](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/val_split20.txt), or 20 out of 1024 total shards) as a 'minival' split, and conduct early stopping based on this 'minival' split. The final accuracy is reported on the original ImageNet validation set. 53 | 54 | We have provided a list of EfficientNet checkpoints:. 55 | 56 | * With baseline ResNet preprocessing, we achieve similar results to the original ICML paper. 57 | * With [AutoAugment](https://arxiv.org/abs/1805.09501) preprocessing, we achieve higher accuracy than the original ICML paper. 58 | * With [RandAugment](https://arxiv.org/abs/1909.13719) preprocessing, accuracy is further improved. 59 | * With [AdvProp](https://arxiv.org/abs/1911.09665), state-of-the-art results (w/o extra data) are achieved. 60 | * With [NoisyStudent](https://arxiv.org/abs/1911.04252), state-of-the-art results (w/ extra JFT-300M unlabeled data) are achieved. 61 | 62 | | | B0 | B1 | B2 | B3 | B4 | B5 | B6 | B7 | B8 | L2-475 | L2 | 63 | |---------- |-------- | ------| ------|------ |------ |------ | --- | --- | --- | --- |--- | 64 | | Baseline preprocessing | 76.7% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b0.tar.gz)) | 78.7% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b1.tar.gz)) | 79.8% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b2.tar.gz)) | 81.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b3.tar.gz)) | 82.5% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b4.tar.gz)) | 83.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b5.tar.gz)) | | || | | | 65 | | AutoAugment (AA) | 77.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b0.tar.gz)) | 79.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b1.tar.gz)) | 80.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b2.tar.gz)) | 81.6% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b3.tar.gz)) | 82.9% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b4.tar.gz)) | 83.6% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b5.tar.gz)) | 84.0% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b6.tar.gz)) | 84.3% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b7.tar.gz)) || | | 66 | | RandAugment (RA) | | | | | | 83.7% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/randaug/efficientnet-b5-randaug.tar.gz)) | | 84.7% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/randaug/efficientnet-b7-randaug.tar.gz)) | | | | 67 | | AdvProp + AA | 77.6% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b0.tar.gz)) | 79.6% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b1.tar.gz)) | 80.5% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b2.tar.gz)) | 81.9% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b3.tar.gz)) | 83.3% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b4.tar.gz)) | 84.3% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b5.tar.gz)) | 84.8% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b6.tar.gz)) | 85.2% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b7.tar.gz)) | 85.5% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b8.tar.gz))|| | | 68 | | NoisyStudent + RA | 78.8% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b0.tar.gz)) | 81.5% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b1.tar.gz)) | 82.4% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b2.tar.gz)) | 84.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b3.tar.gz)) | 85.3% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b4.tar.gz)) | 86.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b5.tar.gz)) | 86.4% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b6.tar.gz)) | 86.9% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b7.tar.gz)) | - |88.2%([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-l2_475.tar.gz))|88.4% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-l2.tar.gz)) | 69 | 70 | 73 | 74 | *To train EfficientNets with AutoAugment ([code](https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py)), simply add option "--augment_name=autoaugment". If you use these checkpoints, you can cite this [paper](https://arxiv.org/abs/1805.09501). 75 | 76 | **To train EfficientNets with RandAugment ([code](https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py)), simply add option "--augment_name=randaugment". For EfficientNet-B5 also add "--randaug_num_layers=2 --randaug_magnitude=17". For EfficientNet-B7 or EfficientNet-B8 also add "--randaug_num_layers=2 --randaug_magnitude=28". If you use these checkpoints, you can cite this [paper](https://arxiv.org/abs/1909.13719). 77 | 78 | * AdvProp training code coming soon. Please set "--advprop_preprocessing=True" for using AdvProp checkpoints. If you use AdvProp checkpoints, you can cite this [paper](https://arxiv.org/abs/1911.09665). 79 | 80 | * NoisyStudent training code coming soon. L2-475 means the same L2 architecture with input image size 475 (Please set "--input_image_size=475" for using this checkpoint). If you use NoisyStudent checkpoints, you can cite this [paper](https://arxiv.org/abs/1911.04252). 81 | 82 | *Note that AdvProp and NoisyStudent performance is derived from baselines that don't use holdout eval set. They will be updated in future." 83 | 84 | A quick way to use these checkpoints is to run: 85 | 86 | $ export MODEL=efficientnet-b0 87 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/${MODEL}.tar.gz 88 | $ tar xf ${MODEL}.tar.gz 89 | $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg 90 | $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.json 91 | $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.json 92 | 93 | Please refer to the following colab for more instructions on how to obtain and use those checkpoints. 94 | 95 | * [`eval_ckpt_example.ipynb`](eval_ckpt_example.ipynb): A colab example to load 96 | EfficientNet pretrained checkpoints files and use the restored model to classify images. 97 | 98 | 99 | ## 3. Using EfficientNet as Feature Extractor 100 | 101 | ``` 102 | import efficientnet_builder 103 | features, endpoints = efficientnet_builder.build_model_base(images, 'efficientnet-b0') 104 | ``` 105 | 106 | * Use `features` for classification finetuning. 107 | * Use `endpoints['reduction_i']` for detection/segmentation, as the last intermediate feature with reduction level `i`. For example, if input image has resolution 224x224, then: 108 | * `endpoints['reduction_1']` has resolution 112x112 109 | * `endpoints['reduction_2']` has resolution 56x56 110 | * `endpoints['reduction_3']` has resolution 28x28 111 | * `endpoints['reduction_4']` has resolution 14x14 112 | * `endpoints['reduction_5']` has resolution 7x7 113 | 114 | ## 4. Training EfficientNets on TPUs. 115 | 116 | 117 | To train this model on Cloud TPU, you will need: 118 | 119 | * A GCE VM instance with an associated Cloud TPU resource 120 | * A GCS bucket to store your training checkpoints (the "model directory") 121 | * Install TensorFlow version >= 1.13 for both GCE VM and Cloud. 122 | 123 | Then train the model: 124 | 125 | $ export PYTHONPATH="$PYTHONPATH:/path/to/models" 126 | $ python main.py --tpu=TPU_NAME --data_dir=DATA_DIR --model_dir=MODEL_DIR 127 | 128 | # TPU_NAME is the name of the TPU node, the same name that appears when you run gcloud compute tpus list, or ctpu ls. 129 | # MODEL_DIR is a GCS location (a URL starting with gs:// where both the GCE VM and the associated Cloud TPU have write access 130 | # DATA_DIR is a GCS location to which both the GCE VM and associated Cloud TPU have read access. 131 | 132 | 133 | For more instructions, please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/efficientnet 134 | -------------------------------------------------------------------------------- /efficientnet_tf/efficientnet_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Model Builder for EfficientNet.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import functools 22 | import os 23 | import re 24 | from absl import logging 25 | import numpy as np 26 | import six 27 | import tensorflow.compat.v1 as tf 28 | 29 | from . import efficientnet_model 30 | from . import utils 31 | MEAN_RGB = [0.485 * 255, 0.456 * 255, 0.406 * 255] 32 | STDDEV_RGB = [0.229 * 255, 0.224 * 255, 0.225 * 255] 33 | 34 | 35 | def efficientnet_params(model_name): 36 | """Get efficientnet params based on model name.""" 37 | params_dict = { 38 | # (width_coefficient, depth_coefficient, resolution, dropout_rate) 39 | 'efficientnet-b0': (1.0, 1.0, 224, 0.2), 40 | 'efficientnet-b1': (1.0, 1.1, 240, 0.2), 41 | 'efficientnet-b2': (1.1, 1.2, 260, 0.3), 42 | 'efficientnet-b3': (1.2, 1.4, 300, 0.3), 43 | 'efficientnet-b4': (1.4, 1.8, 380, 0.4), 44 | 'efficientnet-b5': (1.6, 2.2, 456, 0.4), 45 | 'efficientnet-b6': (1.8, 2.6, 528, 0.5), 46 | 'efficientnet-b7': (2.0, 3.1, 600, 0.5), 47 | 'efficientnet-b8': (2.2, 3.6, 672, 0.5), 48 | 'efficientnet-l2': (4.3, 5.3, 800, 0.5), 49 | } 50 | return params_dict[model_name] 51 | 52 | 53 | class BlockDecoder(object): 54 | """Block Decoder for readability.""" 55 | 56 | def _decode_block_string(self, block_string): 57 | """Gets a block through a string notation of arguments.""" 58 | if six.PY2: 59 | assert isinstance(block_string, (str, unicode)) 60 | else: 61 | assert isinstance(block_string, str) 62 | ops = block_string.split('_') 63 | options = {} 64 | for op in ops: 65 | splits = re.split(r'(\d.*)', op) 66 | if len(splits) >= 2: 67 | key, value = splits[:2] 68 | options[key] = value 69 | 70 | if 's' not in options or len(options['s']) != 2: 71 | raise ValueError('Strides options should be a pair of integers.') 72 | 73 | return efficientnet_model.BlockArgs( 74 | kernel_size=int(options['k']), 75 | num_repeat=int(options['r']), 76 | input_filters=int(options['i']), 77 | output_filters=int(options['o']), 78 | expand_ratio=int(options['e']), 79 | id_skip=('noskip' not in block_string), 80 | se_ratio=float(options['se']) if 'se' in options else None, 81 | strides=[int(options['s'][0]), 82 | int(options['s'][1])], 83 | conv_type=int(options['c']) if 'c' in options else 0, 84 | fused_conv=int(options['f']) if 'f' in options else 0, 85 | space2depth=int(options['d']) if 'd' in options else 0, 86 | condconv=('cc' in block_string), 87 | activation_fn=(tf.nn.relu if int(options['a']) == 0 88 | else tf.nn.swish) if 'a' in options else None) 89 | 90 | def _encode_block_string(self, block): 91 | """Encodes a block to a string.""" 92 | args = [ 93 | 'r%d' % block.num_repeat, 94 | 'k%d' % block.kernel_size, 95 | 's%d%d' % (block.strides[0], block.strides[1]), 96 | 'e%s' % block.expand_ratio, 97 | 'i%d' % block.input_filters, 98 | 'o%d' % block.output_filters, 99 | 'c%d' % block.conv_type, 100 | 'f%d' % block.fused_conv, 101 | 'd%d' % block.space2depth, 102 | ] 103 | if block.se_ratio > 0 and block.se_ratio <= 1: 104 | args.append('se%s' % block.se_ratio) 105 | if block.id_skip is False: # pylint: disable=g-bool-id-comparison 106 | args.append('noskip') 107 | if block.condconv: 108 | args.append('cc') 109 | return '_'.join(args) 110 | 111 | def decode(self, string_list): 112 | """Decodes a list of string notations to specify blocks inside the network. 113 | 114 | Args: 115 | string_list: a list of strings, each string is a notation of block. 116 | 117 | Returns: 118 | A list of namedtuples to represent blocks arguments. 119 | """ 120 | assert isinstance(string_list, list) 121 | blocks_args = [] 122 | for block_string in string_list: 123 | blocks_args.append(self._decode_block_string(block_string)) 124 | return blocks_args 125 | 126 | def encode(self, blocks_args): 127 | """Encodes a list of Blocks to a list of strings. 128 | 129 | Args: 130 | blocks_args: A list of namedtuples to represent blocks arguments. 131 | Returns: 132 | a list of strings, each string is a notation of block. 133 | """ 134 | block_strings = [] 135 | for block in blocks_args: 136 | block_strings.append(self._encode_block_string(block)) 137 | return block_strings 138 | 139 | 140 | def swish(features, use_native=True, use_hard=False): 141 | """Computes the Swish activation function. 142 | 143 | We provide three alternnatives: 144 | - Native tf.nn.swish, use less memory during training than composable swish. 145 | - Quantization friendly hard swish. 146 | - A composable swish, equivalant to tf.nn.swish, but more general for 147 | finetuning and TF-Hub. 148 | 149 | Args: 150 | features: A `Tensor` representing preactivation values. 151 | use_native: Whether to use the native swish from tf.nn that uses a custom 152 | gradient to reduce memory usage, or to use customized swish that uses 153 | default TensorFlow gradient computation. 154 | use_hard: Whether to use quantization-friendly hard swish. 155 | 156 | Returns: 157 | The activation value. 158 | """ 159 | if use_native and use_hard: 160 | raise ValueError('Cannot specify both use_native and use_hard.') 161 | 162 | if use_native: 163 | return tf.nn.swish(features) 164 | 165 | if use_hard: 166 | return features * tf.nn.relu6(features + np.float32(3)) * (1. / 6.) 167 | 168 | features = tf.convert_to_tensor(features, name='features') 169 | return features * tf.nn.sigmoid(features) 170 | 171 | 172 | _DEFAULT_BLOCKS_ARGS = [ 173 | 'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25', 174 | 'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25', 175 | 'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25', 176 | 'r1_k3_s11_e6_i192_o320_se0.25', 177 | ] 178 | 179 | 180 | def efficientnet(width_coefficient=None, 181 | depth_coefficient=None, 182 | dropout_rate=0.2, 183 | survival_prob=0.8): 184 | """Creates a efficientnet model.""" 185 | global_params = efficientnet_model.GlobalParams( 186 | blocks_args=_DEFAULT_BLOCKS_ARGS, 187 | batch_norm_momentum=0.99, 188 | batch_norm_epsilon=1e-3, 189 | dropout_rate=dropout_rate, 190 | survival_prob=survival_prob, 191 | data_format='channels_last', 192 | num_classes=1000, 193 | width_coefficient=width_coefficient, 194 | depth_coefficient=depth_coefficient, 195 | depth_divisor=8, 196 | min_depth=None, 197 | relu_fn=tf.nn.swish, 198 | # The default is TPU-specific batch norm. 199 | # The alternative is tf.layers.BatchNormalization. 200 | batch_norm=utils.train_batch_norm, # TPU-specific requirement. 201 | use_se=True, 202 | clip_projection_output=False) 203 | return global_params 204 | 205 | 206 | def get_model_params(model_name, override_params): 207 | """Get the block args and global params for a given model.""" 208 | if model_name.startswith('efficientnet'): 209 | width_coefficient, depth_coefficient, _, dropout_rate = ( 210 | efficientnet_params(model_name)) 211 | global_params = efficientnet( 212 | width_coefficient, depth_coefficient, dropout_rate) 213 | else: 214 | raise NotImplementedError('model name is not pre-defined: %s' % model_name) 215 | 216 | if override_params: 217 | # ValueError will be raised here if override_params has fields not included 218 | # in global_params. 219 | global_params = global_params._replace(**override_params) 220 | 221 | decoder = BlockDecoder() 222 | blocks_args = decoder.decode(global_params.blocks_args) 223 | 224 | logging.info('global_params= %s', global_params) 225 | return blocks_args, global_params 226 | 227 | 228 | def build_model(images, 229 | model_name, 230 | training, 231 | override_params=None, 232 | model_dir=None, 233 | fine_tuning=False, 234 | features_only=False, 235 | pooled_features_only=False): 236 | """A helper function to create a model and return predicted logits. 237 | 238 | Args: 239 | images: input images tensor. 240 | model_name: string, the predefined model name. 241 | training: boolean, whether the model is constructed for training. 242 | override_params: A dictionary of params for overriding. Fields must exist in 243 | efficientnet_model.GlobalParams. 244 | model_dir: string, optional model dir for saving configs. 245 | fine_tuning: boolean, whether the model is used for finetuning. 246 | features_only: build the base feature network only (excluding final 247 | 1x1 conv layer, global pooling, dropout and fc head). 248 | pooled_features_only: build the base network for features extraction (after 249 | 1x1 conv layer and global pooling, but before dropout and fc head). 250 | 251 | Returns: 252 | logits: the logits tensor of classes. 253 | endpoints: the endpoints for each layer. 254 | 255 | Raises: 256 | When model_name specified an undefined model, raises NotImplementedError. 257 | When override_params has invalid fields, raises ValueError. 258 | """ 259 | assert isinstance(images, tf.Tensor) 260 | assert not (features_only and pooled_features_only) 261 | 262 | # For backward compatibility. 263 | if override_params and override_params.get('drop_connect_rate', None): 264 | override_params['survival_prob'] = 1 - override_params['drop_connect_rate'] 265 | 266 | if not training or fine_tuning: 267 | if not override_params: 268 | override_params = {} 269 | override_params['batch_norm'] = utils.eval_batch_norm 270 | if fine_tuning: 271 | override_params['relu_fn'] = functools.partial(swish, use_native=False) 272 | blocks_args, global_params = get_model_params(model_name, override_params) 273 | 274 | if model_dir: 275 | param_file = os.path.join(model_dir, 'model_params.txt') 276 | if not tf.gfile.Exists(param_file): 277 | if not tf.gfile.Exists(model_dir): 278 | tf.gfile.MakeDirs(model_dir) 279 | with tf.gfile.GFile(param_file, 'w') as f: 280 | logging.info('writing to %s', param_file) 281 | f.write('model_name= %s\n\n' % model_name) 282 | f.write('global_params= %s\n\n' % str(global_params)) 283 | f.write('blocks_args= %s\n\n' % str(blocks_args)) 284 | 285 | with tf.variable_scope(model_name): 286 | model = efficientnet_model.Model(blocks_args, global_params) 287 | outputs = model( 288 | images, 289 | training=training, 290 | features_only=features_only, 291 | pooled_features_only=pooled_features_only) 292 | if features_only: 293 | outputs = tf.identity(outputs, 'features') 294 | elif pooled_features_only: 295 | outputs = tf.identity(outputs, 'pooled_features') 296 | else: 297 | outputs = tf.identity(outputs, 'logits') 298 | return outputs, model.endpoints 299 | 300 | 301 | def build_model_base(images, model_name, training, override_params=None): 302 | """Create a base feature network and return the features before pooling. 303 | 304 | Args: 305 | images: input images tensor. 306 | model_name: string, the predefined model name. 307 | training: boolean, whether the model is constructed for training. 308 | override_params: A dictionary of params for overriding. Fields must exist in 309 | efficientnet_model.GlobalParams. 310 | 311 | Returns: 312 | features: base features before pooling. 313 | endpoints: the endpoints for each layer. 314 | 315 | Raises: 316 | When model_name specified an undefined model, raises NotImplementedError. 317 | When override_params has invalid fields, raises ValueError. 318 | """ 319 | assert isinstance(images, tf.Tensor) 320 | # For backward compatibility. 321 | if override_params and override_params.get('drop_connect_rate', None): 322 | override_params['survival_prob'] = 1 - override_params['drop_connect_rate'] 323 | 324 | blocks_args, global_params = get_model_params(model_name, override_params) 325 | 326 | with tf.variable_scope(model_name): 327 | model = efficientnet_model.Model(blocks_args, global_params) 328 | features = model(images, training=training, features_only=True) 329 | 330 | features = tf.identity(features, 'features') 331 | return features, model.endpoints 332 | -------------------------------------------------------------------------------- /efficientnet_tf/condconv/condconv_layers.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """CondConv implementations in Tensorflow Layers. 16 | 17 | [1] Brandon Yang, Gabriel Bender, Quoc V. Le, Jiquan Ngiam 18 | CondConv: Conditionally Parameterized Convolutions for Efficient Inference. 19 | NeurIPS'19, https://arxiv.org/abs/1904.04971 20 | """ 21 | 22 | from __future__ import absolute_import 23 | from __future__ import division 24 | #Standard imports 25 | from __future__ import print_function 26 | 27 | import numpy as np 28 | import tensorflow.compat.v1 as tf 29 | 30 | 31 | def get_condconv_initializer(initializer, num_experts, expert_shape): 32 | """Wraps the initializer to correctly initialize CondConv variables. 33 | 34 | CondConv initializes biases and kernels in a num_experts x num_params 35 | matrix for efficient computation. This wrapper ensures that each expert 36 | is correctly initialized with the given initializer before being flattened 37 | into the correctly shaped CondConv variable. 38 | 39 | Arguments: 40 | initializer: The initializer to apply for each individual expert. 41 | num_experts: The number of experts to be initialized. 42 | expert_shape: The original shape of each individual expert. 43 | 44 | Returns: 45 | The initializer for the num_experts x num_params CondConv variable. 46 | """ 47 | def condconv_initializer(expected_shape, dtype=None, partition=None): 48 | """CondConv initializer function.""" 49 | num_params = np.prod(expert_shape) 50 | if (len(expected_shape) != 2 or expected_shape[0] != num_experts or 51 | expected_shape[1] != num_params): 52 | raise (ValueError( 53 | 'CondConv variables must have shape [num_experts, num_params]')) 54 | flattened_kernels = [] 55 | for _ in range(num_experts): 56 | kernel = initializer(expert_shape, dtype, partition) 57 | flattened_kernels.append(tf.reshape(kernel, [-1])) 58 | return tf.stack(flattened_kernels) 59 | 60 | return condconv_initializer 61 | 62 | 63 | class CondConv2D(tf.keras.layers.Conv2D): 64 | """2D conditional convolution layer (e.g. spatial convolution over images). 65 | 66 | Attributes: 67 | filters: Integer, the dimensionality of the output space (i.e. the number of 68 | output filters in the convolution). 69 | kernel_size: An integer or tuple/list of 2 integers, specifying the height 70 | and width of the 2D convolution window. Can be a single integer to specify 71 | the same value for all spatial dimensions. 72 | num_experts: The number of expert kernels and biases in the CondConv layer. 73 | strides: An integer or tuple/list of 2 integers, specifying the strides of 74 | the convolution along the height and width. Can be a single integer to 75 | specify the same value for all spatial dimensions. Specifying any stride 76 | value != 1 is incompatible with specifying any `dilation_rate` value != 1. 77 | padding: one of `"valid"` or `"same"` (case-insensitive). 78 | data_format: A string, one of `channels_last` (default) or `channels_first`. 79 | The ordering of the dimensions in the inputs. `channels_last` corresponds 80 | to inputs with shape `(batch, height, width, channels)` while 81 | `channels_first` corresponds to inputs with shape `(batch, channels, 82 | height, width)`. It defaults to the `image_data_format` value found in 83 | your Keras config file at `~/.keras/keras.json`. If you never set it, then 84 | it will be "channels_last". 85 | dilation_rate: an integer or tuple/list of 2 integers, specifying the 86 | dilation rate to use for dilated convolution. Can be a single integer to 87 | specify the same value for all spatial dimensions. Currently, specifying 88 | any `dilation_rate` value != 1 is incompatible with specifying any stride 89 | value != 1. 90 | activation: Activation function to use. If you don't specify anything, no 91 | activation is applied 92 | (ie. "linear" activation: `a(x) = x`). 93 | use_bias: Boolean, whether the layer uses a bias vector. 94 | kernel_initializer: Initializer for the `kernel` weights matrix. 95 | bias_initializer: Initializer for the bias vector. 96 | kernel_regularizer: Regularizer function applied to the `kernel` weights 97 | matrix. 98 | bias_regularizer: Regularizer function applied to the bias vector. 99 | activity_regularizer: Regularizer function applied to the output of the 100 | layer (its "activation").. 101 | kernel_constraint: Constraint function applied to the kernel matrix. 102 | bias_constraint: Constraint function applied to the bias vector. 103 | Input shape: 104 | 4D tensor with shape: `(samples, channels, rows, cols)` if 105 | data_format='channels_first' 106 | or 4D tensor with shape: `(samples, rows, cols, channels)` if 107 | data_format='channels_last'. 108 | Output shape: 109 | 4D tensor with shape: `(samples, filters, new_rows, new_cols)` if 110 | data_format='channels_first' 111 | or 4D tensor with shape: `(samples, new_rows, new_cols, filters)` if 112 | data_format='channels_last'. `rows` and `cols` values might have changed 113 | due to padding. 114 | """ 115 | 116 | def __init__(self, 117 | filters, 118 | kernel_size, 119 | num_experts, 120 | strides=(1, 1), 121 | padding='valid', 122 | data_format=None, 123 | dilation_rate=(1, 1), 124 | activation=None, 125 | use_bias=True, 126 | kernel_initializer='glorot_uniform', 127 | bias_initializer='zeros', 128 | kernel_regularizer=None, 129 | bias_regularizer=None, 130 | activity_regularizer=None, 131 | kernel_constraint=None, 132 | bias_constraint=None, 133 | **kwargs): 134 | super(CondConv2D, self).__init__( 135 | filters=filters, 136 | kernel_size=kernel_size, 137 | strides=strides, 138 | padding=padding, 139 | data_format=data_format, 140 | dilation_rate=dilation_rate, 141 | activation=activation, 142 | use_bias=use_bias, 143 | kernel_initializer=kernel_initializer, 144 | bias_initializer=bias_initializer, 145 | kernel_regularizer=kernel_regularizer, 146 | bias_regularizer=bias_regularizer, 147 | activity_regularizer=activity_regularizer, 148 | kernel_constraint=kernel_constraint, 149 | bias_constraint=bias_constraint, 150 | **kwargs) 151 | if num_experts < 1: 152 | raise ValueError('A CondConv layer must have at least one expert.') 153 | self.num_experts = num_experts 154 | if self.data_format == 'channels_first': 155 | self.converted_data_format = 'NCHW' 156 | else: 157 | self.converted_data_format = 'NHWC' 158 | 159 | def build(self, input_shape): 160 | if len(input_shape) != 4: 161 | raise ValueError( 162 | 'Inputs to `CondConv2D` should have rank 4. ' 163 | 'Received input shape:', str(input_shape)) 164 | input_shape = tf.TensorShape(input_shape) 165 | channel_axis = self._get_channel_axis() 166 | if input_shape.dims[channel_axis].value is None: 167 | raise ValueError('The channel dimension of the inputs ' 168 | 'should be defined. Found `None`.') 169 | input_dim = int(input_shape[channel_axis]) 170 | 171 | self.kernel_shape = self.kernel_size + (input_dim, self.filters) 172 | kernel_num_params = 1 173 | for kernel_dim in self.kernel_shape: 174 | kernel_num_params *= kernel_dim 175 | condconv_kernel_shape = (self.num_experts, kernel_num_params) 176 | self.condconv_kernel = self.add_weight( 177 | name='condconv_kernel', 178 | shape=condconv_kernel_shape, 179 | initializer=get_condconv_initializer(self.kernel_initializer, 180 | self.num_experts, 181 | self.kernel_shape), 182 | regularizer=self.kernel_regularizer, 183 | constraint=self.kernel_constraint, 184 | trainable=True, 185 | dtype=self.dtype) 186 | 187 | if self.use_bias: 188 | self.bias_shape = (self.filters,) 189 | condconv_bias_shape = (self.num_experts, self.filters) 190 | self.condconv_bias = self.add_weight( 191 | name='condconv_bias', 192 | shape=condconv_bias_shape, 193 | initializer=get_condconv_initializer(self.bias_initializer, 194 | self.num_experts, 195 | self.bias_shape), 196 | regularizer=self.bias_regularizer, 197 | constraint=self.bias_constraint, 198 | trainable=True, 199 | dtype=self.dtype) 200 | else: 201 | self.bias = None 202 | 203 | self.input_spec = tf.layers.InputSpec( 204 | ndim=self.rank + 2, axes={channel_axis: input_dim}) 205 | 206 | self.built = True 207 | 208 | def call(self, inputs, routing_weights): 209 | # Compute example dependent kernels 210 | kernels = tf.matmul(routing_weights, self.condconv_kernel) 211 | batch_size = inputs.shape[0].value 212 | inputs = tf.split(inputs, batch_size, 0) 213 | kernels = tf.split(kernels, batch_size, 0) 214 | # Apply example-dependent convolution to each example in the batch 215 | outputs_list = [] 216 | for input_tensor, kernel in zip(inputs, kernels): 217 | kernel = tf.reshape(kernel, self.kernel_shape) 218 | outputs_list.append( 219 | tf.nn.convolution( 220 | input_tensor, 221 | kernel, 222 | strides=self.strides, 223 | padding=self._get_padding_op(), 224 | dilations=self.dilation_rate, 225 | data_format=self.converted_data_format)) 226 | outputs = tf.concat(outputs_list, 0) 227 | 228 | if self.use_bias: 229 | # Compute example-dependent biases 230 | biases = tf.matmul(routing_weights, self.condconv_bias) 231 | outputs = tf.split(outputs, batch_size, 0) 232 | biases = tf.split(biases, batch_size, 0) 233 | # Add example-dependent bias to each example in the batch 234 | bias_outputs_list = [] 235 | for output, bias in zip(outputs, biases): 236 | bias = tf.squeeze(bias, axis=0) 237 | bias_outputs_list.append( 238 | tf.nn.bias_add(output, bias, 239 | data_format=self.converted_data_format)) 240 | outputs = tf.concat(bias_outputs_list, 0) 241 | 242 | if self.activation is not None: 243 | return self.activation(outputs) 244 | return outputs 245 | 246 | def get_config(self): 247 | config = {'num_experts': self.num_experts} 248 | base_config = super(CondConv2D, self).get_config() 249 | return dict(list(base_config.items()) + list(config.items())) 250 | 251 | def _get_channel_axis(self): 252 | if self.data_format == 'channels_first': 253 | return 1 254 | else: 255 | return -1 256 | 257 | def _get_padding_op(self): 258 | if self.padding == 'causal': 259 | op_padding = 'valid' 260 | else: 261 | op_padding = self.padding 262 | if not isinstance(op_padding, (list, tuple)): 263 | op_padding = op_padding.upper() 264 | return op_padding 265 | 266 | 267 | class DepthwiseCondConv2D(tf.keras.layers.DepthwiseConv2D): 268 | """Depthwise separable 2D conditional convolution layer. 269 | 270 | This layer extends the base depthwise 2D convolution layer to compute 271 | example-dependent parameters. A DepthwiseCondConv2D layer has 'num_experts` 272 | kernels and biases. It computes a kernel and bias for each example as a 273 | weighted sum of experts using the input example-dependent routing weights, 274 | then applies the depthwise convolution to each example. 275 | 276 | Attributes: 277 | kernel_size: An integer or tuple/list of 2 integers, specifying the height 278 | and width of the 2D convolution window. Can be a single integer to specify 279 | the same value for all spatial dimensions. 280 | num_experts: The number of expert kernels and biases in the 281 | DepthwiseCondConv2D layer. 282 | strides: An integer or tuple/list of 2 integers, specifying the strides of 283 | the convolution along the height and width. Can be a single integer to 284 | specify the same value for all spatial dimensions. Specifying any stride 285 | value != 1 is incompatible with specifying any `dilation_rate` value != 1. 286 | padding: one of `'valid'` or `'same'` (case-insensitive). 287 | depth_multiplier: The number of depthwise convolution output channels for 288 | each input channel. The total number of depthwise convolution output 289 | channels will be equal to `filters_in * depth_multiplier`. 290 | data_format: A string, one of `channels_last` (default) or `channels_first`. 291 | The ordering of the dimensions in the inputs. `channels_last` corresponds 292 | to inputs with shape `(batch, height, width, channels)` while 293 | `channels_first` corresponds to inputs with shape `(batch, channels, 294 | height, width)`. It defaults to the `image_data_format` value found in 295 | your Keras config file at `~/.keras/keras.json`. If you never set it, then 296 | it will be 'channels_last'. 297 | activation: Activation function to use. If you don't specify anything, no 298 | activation is applied 299 | (ie. 'linear' activation: `a(x) = x`). 300 | use_bias: Boolean, whether the layer uses a bias vector. 301 | depthwise_initializer: Initializer for the depthwise kernel matrix. 302 | bias_initializer: Initializer for the bias vector. 303 | depthwise_regularizer: Regularizer function applied to the depthwise kernel 304 | matrix. 305 | bias_regularizer: Regularizer function applied to the bias vector. 306 | activity_regularizer: Regularizer function applied to the output of the 307 | layer (its 'activation'). 308 | depthwise_constraint: Constraint function applied to the depthwise kernel 309 | matrix. 310 | bias_constraint: Constraint function applied to the bias vector. 311 | Input shape: 312 | 4D tensor with shape: `[batch, channels, rows, cols]` if 313 | data_format='channels_first' 314 | or 4D tensor with shape: `[batch, rows, cols, channels]` if 315 | data_format='channels_last'. 316 | Output shape: 317 | 4D tensor with shape: `[batch, filters, new_rows, new_cols]` if 318 | data_format='channels_first' 319 | or 4D tensor with shape: `[batch, new_rows, new_cols, filters]` if 320 | data_format='channels_last'. `rows` and `cols` values might have changed 321 | due to padding. 322 | """ 323 | 324 | def __init__(self, 325 | kernel_size, 326 | num_experts, 327 | strides=(1, 1), 328 | padding='valid', 329 | depth_multiplier=1, 330 | data_format=None, 331 | activation=None, 332 | use_bias=True, 333 | depthwise_initializer='glorot_uniform', 334 | bias_initializer='zeros', 335 | depthwise_regularizer=None, 336 | bias_regularizer=None, 337 | activity_regularizer=None, 338 | depthwise_constraint=None, 339 | bias_constraint=None, 340 | **kwargs): 341 | super(DepthwiseCondConv2D, self).__init__( 342 | kernel_size=kernel_size, 343 | strides=strides, 344 | padding=padding, 345 | depth_multiplier=depth_multiplier, 346 | data_format=data_format, 347 | activation=activation, 348 | use_bias=use_bias, 349 | depthwise_initializer=depthwise_initializer, 350 | bias_initializer=bias_initializer, 351 | depthwise_regularizer=depthwise_regularizer, 352 | bias_regularizer=bias_regularizer, 353 | activity_regularizer=activity_regularizer, 354 | depthwise_constraint=depthwise_constraint, 355 | bias_constraint=bias_constraint, 356 | **kwargs) 357 | if num_experts < 1: 358 | raise ValueError('A CondConv layer must have at least one expert.') 359 | self.num_experts = num_experts 360 | if self.data_format == 'channels_first': 361 | self.converted_data_format = 'NCHW' 362 | else: 363 | self.converted_data_format = 'NHWC' 364 | 365 | def build(self, input_shape): 366 | if len(input_shape) < 4: 367 | raise ValueError( 368 | 'Inputs to `DepthwiseCondConv2D` should have rank 4. ' 369 | 'Received input shape:', str(input_shape)) 370 | input_shape = tf.TensorShape(input_shape) 371 | if self.data_format == 'channels_first': 372 | channel_axis = 1 373 | else: 374 | channel_axis = 3 375 | if input_shape.dims[channel_axis].value is None: 376 | raise ValueError('The channel dimension of the inputs to ' 377 | '`DepthwiseConv2D` ' 378 | 'should be defined. Found `None`.') 379 | input_dim = int(input_shape[channel_axis]) 380 | self.depthwise_kernel_shape = (self.kernel_size[0], self.kernel_size[1], 381 | input_dim, self.depth_multiplier) 382 | 383 | depthwise_kernel_num_params = 1 384 | for dim in self.depthwise_kernel_shape: 385 | depthwise_kernel_num_params *= dim 386 | depthwise_condconv_kernel_shape = (self.num_experts, 387 | depthwise_kernel_num_params) 388 | 389 | self.depthwise_condconv_kernel = self.add_weight( 390 | shape=depthwise_condconv_kernel_shape, 391 | initializer=get_condconv_initializer(self.depthwise_initializer, 392 | self.num_experts, 393 | self.depthwise_kernel_shape), 394 | name='depthwise_condconv_kernel', 395 | regularizer=self.depthwise_regularizer, 396 | constraint=self.depthwise_constraint, 397 | trainable=True) 398 | 399 | if self.use_bias: 400 | bias_dim = input_dim * self.depth_multiplier 401 | self.bias_shape = (bias_dim,) 402 | condconv_bias_shape = (self.num_experts, bias_dim) 403 | self.condconv_bias = self.add_weight( 404 | name='condconv_bias', 405 | shape=condconv_bias_shape, 406 | initializer=get_condconv_initializer(self.bias_initializer, 407 | self.num_experts, 408 | self.bias_shape), 409 | regularizer=self.bias_regularizer, 410 | constraint=self.bias_constraint, 411 | trainable=True, 412 | dtype=self.dtype) 413 | else: 414 | self.bias = None 415 | # Set input spec. 416 | self.input_spec = tf.layers.InputSpec( 417 | ndim=4, axes={channel_axis: input_dim}) 418 | self.built = True 419 | 420 | def call(self, inputs, routing_weights): 421 | # Compute example dependent depthwise kernels 422 | depthwise_kernels = tf.matmul(routing_weights, 423 | self.depthwise_condconv_kernel) 424 | batch_size = inputs.shape[0].value 425 | inputs = tf.split(inputs, batch_size, 0) 426 | depthwise_kernels = tf.split(depthwise_kernels, batch_size, 0) 427 | # Apply example-dependent depthwise convolution to each example in the batch 428 | outputs_list = [] 429 | for input_tensor, depthwise_kernel in zip(inputs, depthwise_kernels): 430 | depthwise_kernel = tf.reshape(depthwise_kernel, 431 | self.depthwise_kernel_shape) 432 | if self.data_format == 'channels_first': 433 | converted_strides = (1, 1) + self.strides 434 | else: 435 | converted_strides = (1,) + self.strides + (1,) 436 | outputs_list.append( 437 | tf.nn.depthwise_conv2d( 438 | input_tensor, 439 | depthwise_kernel, 440 | strides=converted_strides, 441 | padding=self.padding.upper(), 442 | dilations=self.dilation_rate, 443 | data_format=self.converted_data_format)) 444 | outputs = tf.concat(outputs_list, 0) 445 | 446 | if self.use_bias: 447 | # Compute example-dependent biases 448 | biases = tf.matmul(routing_weights, self.condconv_bias) 449 | outputs = tf.split(outputs, batch_size, 0) 450 | biases = tf.split(biases, batch_size, 0) 451 | # Add example-dependent bias to each example in the batch 452 | bias_outputs_list = [] 453 | for output, bias in zip(outputs, biases): 454 | bias = tf.squeeze(bias, axis=0) 455 | bias_outputs_list.append( 456 | tf.nn.bias_add(output, bias, 457 | data_format=self.converted_data_format)) 458 | outputs = tf.concat(bias_outputs_list, 0) 459 | 460 | if self.activation is not None: 461 | return self.activation(outputs) 462 | 463 | return outputs 464 | 465 | def get_config(self): 466 | config = {'num_experts': self.num_experts} 467 | base_config = super(DepthwiseCondConv2D, self).get_config() 468 | return dict(list(base_config.items()) + list(config.items())) 469 | -------------------------------------------------------------------------------- /efficientnet_tf/imagenet_input.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Efficient ImageNet input pipeline using tf.data.Dataset.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import abc 22 | import collections 23 | import functools 24 | import os 25 | 26 | from absl import logging 27 | import six 28 | import tensorflow.compat.v1 as tf 29 | 30 | import preprocessing 31 | 32 | 33 | def build_image_serving_input_fn(image_size, 34 | batch_size=None, 35 | resize_method=None): 36 | """Builds a serving input fn for raw images.""" 37 | 38 | def _image_serving_input_fn(): 39 | """Serving input fn for raw images.""" 40 | 41 | def _preprocess_image(image_bytes): 42 | """Preprocess a single raw image.""" 43 | image = preprocessing.preprocess_image( 44 | image_bytes=image_bytes, 45 | is_training=False, 46 | image_size=image_size, 47 | resize_method=resize_method) 48 | return image 49 | 50 | image_bytes_list = tf.placeholder( 51 | shape=[batch_size], 52 | dtype=tf.string, 53 | ) 54 | images = tf.map_fn( 55 | _preprocess_image, image_bytes_list, back_prop=False, dtype=tf.float32) 56 | return tf.estimator.export.ServingInputReceiver( 57 | images, {'image_bytes': image_bytes_list}) 58 | return _image_serving_input_fn 59 | 60 | 61 | class ImageNetTFExampleInput(six.with_metaclass(abc.ABCMeta, object)): 62 | """Base class for ImageNet input_fn generator.""" 63 | 64 | def __init__(self, 65 | is_training, 66 | use_bfloat16, 67 | num_cores=8, 68 | image_size=224, 69 | transpose_input=False, 70 | num_label_classes=1000, 71 | include_background_label=False, 72 | augment_name=None, 73 | mixup_alpha=0.0, 74 | randaug_num_layers=None, 75 | randaug_magnitude=None, 76 | resize_method=None): 77 | """Constructor. 78 | 79 | Args: 80 | is_training: `bool` for whether the input is for training 81 | use_bfloat16: If True, use bfloat16 precision; else use float32. 82 | num_cores: `int` for the number of TPU cores 83 | image_size: `int` for image size (both width and height). 84 | transpose_input: 'bool' for whether to use the double transpose trick 85 | num_label_classes: number of label classes. Default to 1000 for ImageNet. 86 | include_background_label: If true, label #0 is reserved for background. 87 | augment_name: `string` that is the name of the augmentation method to 88 | apply to the image. `autoaugment` if AutoAugment is to be used or 89 | `randaugment` if RandAugment is to be used. If the value is `None` no no 90 | augmentation method will be applied applied. See autoaugment.py for more 91 | details. 92 | mixup_alpha: float to control the strength of Mixup regularization, set to 93 | 0.0 to disable. 94 | randaug_num_layers: 'int', if RandAug is used, what should the number of 95 | layers be. See autoaugment.py for detailed description. 96 | randaug_magnitude: 'int', if RandAug is used, what should the magnitude 97 | be. See autoaugment.py for detailed description. 98 | resize_method: If None, use bicubic in default. 99 | """ 100 | self.image_preprocessing_fn = preprocessing.preprocess_image 101 | self.is_training = is_training 102 | self.use_bfloat16 = use_bfloat16 103 | self.num_cores = num_cores 104 | self.transpose_input = transpose_input 105 | self.image_size = image_size 106 | self.include_background_label = include_background_label 107 | self.num_label_classes = num_label_classes 108 | if include_background_label: 109 | self.num_label_classes += 1 110 | self.augment_name = augment_name 111 | self.mixup_alpha = mixup_alpha 112 | self.randaug_num_layers = randaug_num_layers 113 | self.randaug_magnitude = randaug_magnitude 114 | self.resize_method = resize_method 115 | 116 | def set_shapes(self, batch_size, images, labels): 117 | """Statically set the batch_size dimension.""" 118 | if self.transpose_input: 119 | images.set_shape(images.get_shape().merge_with( 120 | tf.TensorShape([None, None, None, batch_size]))) 121 | labels.set_shape(labels.get_shape().merge_with( 122 | tf.TensorShape([batch_size, None]))) 123 | # Convert to R1 tensors for fast transfer to device. 124 | images = tf.reshape(images, [-1]) 125 | else: 126 | images.set_shape(images.get_shape().merge_with( 127 | tf.TensorShape([batch_size, None, None, None]))) 128 | labels.set_shape(labels.get_shape().merge_with( 129 | tf.TensorShape([batch_size, None]))) 130 | 131 | return images, labels 132 | 133 | def mixup(self, batch_size, alpha, images, labels): 134 | """Applies Mixup regularization to a batch of images and labels. 135 | 136 | [1] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, David Lopez-Paz 137 | Mixup: Beyond Empirical Risk Minimization. 138 | ICLR'18, https://arxiv.org/abs/1710.09412 139 | 140 | Arguments: 141 | batch_size: The input batch size for images and labels. 142 | alpha: Float that controls the strength of Mixup regularization. 143 | images: A batch of images of shape [batch_size, ...] 144 | labels: A batch of labels of shape [batch_size, num_classes] 145 | 146 | Returns: 147 | A tuple of (images, labels) with the same dimensions as the input with 148 | Mixup regularization applied. 149 | """ 150 | mix_weight = tf.distributions.Beta(alpha, alpha).sample([batch_size, 1]) 151 | mix_weight = tf.maximum(mix_weight, 1. - mix_weight) 152 | images_mix_weight = tf.cast( 153 | tf.reshape(mix_weight, [batch_size, 1, 1, 1]), images.dtype) 154 | # Mixup on a single batch is implemented by taking a weighted sum with the 155 | # same batch in reverse. 156 | images_mix = ( 157 | images * images_mix_weight + images[::-1] * (1. - images_mix_weight)) 158 | labels_mix = labels * mix_weight + labels[::-1] * (1. - mix_weight) 159 | return images_mix, labels_mix 160 | 161 | def dataset_parser(self, value): 162 | """Parses an image and its label from a serialized ResNet-50 TFExample. 163 | 164 | Args: 165 | value: serialized string containing an ImageNet TFExample. 166 | 167 | Returns: 168 | Returns a tuple of (image, label) from the TFExample. 169 | """ 170 | keys_to_features = { 171 | 'image/encoded': tf.FixedLenFeature((), tf.string, ''), 172 | 'image/class/label': tf.FixedLenFeature([], tf.int64, -1), 173 | } 174 | 175 | parsed = tf.parse_single_example(value, keys_to_features) 176 | image_bytes = tf.reshape(parsed['image/encoded'], shape=[]) 177 | 178 | image = self.image_preprocessing_fn( 179 | image_bytes=image_bytes, 180 | is_training=self.is_training, 181 | image_size=self.image_size, 182 | use_bfloat16=self.use_bfloat16, 183 | augment_name=self.augment_name, 184 | randaug_num_layers=self.randaug_num_layers, 185 | randaug_magnitude=self.randaug_magnitude, 186 | resize_method=self.resize_method) 187 | 188 | # The labels will be in range [1,1000], 0 is reserved for background 189 | label = tf.cast( 190 | tf.reshape(parsed['image/class/label'], shape=[]), dtype=tf.int32) 191 | 192 | if not self.include_background_label: 193 | # Subtract 1 if the background label is discarded. 194 | label -= 1 195 | 196 | onehot_label = tf.one_hot(label, self.num_label_classes) 197 | 198 | return image, onehot_label 199 | 200 | @abc.abstractmethod 201 | def make_source_dataset(self, index, num_hosts): 202 | """Makes dataset of serialized TFExamples. 203 | 204 | The returned dataset will contain `tf.string` tensors, but these strings are 205 | serialized `TFExample` records that will be parsed by `dataset_parser`. 206 | 207 | If self.is_training, the dataset should be infinite. 208 | 209 | Args: 210 | index: current host index. 211 | num_hosts: total number of hosts. 212 | 213 | Returns: 214 | A `tf.data.Dataset` object. 215 | """ 216 | return 217 | 218 | def input_fn(self, params): 219 | """Input function which provides a single batch for train or eval. 220 | 221 | Args: 222 | params: `dict` of parameters passed from the `TPUEstimator`. 223 | `params['batch_size']` is always provided and should be used as the 224 | effective batch size. 225 | 226 | Returns: 227 | A `tf.data.Dataset` object. 228 | """ 229 | # Retrieves the batch size for the current shard. The # of shards is 230 | # computed according to the input pipeline deployment. See 231 | # tf.estimator.tpu.RunConfig for details. 232 | batch_size = params['batch_size'] 233 | 234 | if 'context' in params: 235 | current_host = params['context'].current_input_fn_deployment()[1] 236 | num_hosts = params['context'].num_hosts 237 | else: 238 | current_host = 0 239 | num_hosts = 1 240 | 241 | dataset = self.make_source_dataset(current_host, num_hosts) 242 | 243 | # Use the fused map-and-batch operation. 244 | # 245 | # For XLA, we must used fixed shapes. Because we repeat the source training 246 | # dataset indefinitely, we can use `drop_remainder=True` to get fixed-size 247 | # batches without dropping any training examples. 248 | # 249 | # When evaluating, `drop_remainder=True` prevents accidentally evaluating 250 | # the same image twice by dropping the final batch if it is less than a full 251 | # batch size. As long as this validation is done with consistent batch size, 252 | # exactly the same images will be used. 253 | dataset = dataset.map(self.dataset_parser, 64).batch(batch_size, True) 254 | 255 | # Apply Mixup 256 | if self.is_training and self.mixup_alpha > 0.0: 257 | dataset = dataset.map( 258 | functools.partial(self.mixup, batch_size, self.mixup_alpha), 259 | num_parallel_calls=64) 260 | 261 | # Transpose for performance on TPU 262 | if self.transpose_input: 263 | dataset = dataset.map( 264 | lambda images, labels: (tf.transpose(images, [1, 2, 3, 0]), labels), 265 | num_parallel_calls=64) 266 | 267 | # Assign static batch size dimension 268 | dataset = dataset.map(functools.partial(self.set_shapes, batch_size), 64) 269 | 270 | # Prefetch overlaps in-feed with training 271 | dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) 272 | options = tf.data.Options() 273 | options.experimental_deterministic = False 274 | options.experimental_threading.max_intra_op_parallelism = 1 275 | options.experimental_threading.private_threadpool_size = 48 276 | dataset = dataset.with_options(options) 277 | 278 | return dataset 279 | 280 | 281 | class ImageNetInput(ImageNetTFExampleInput): 282 | """Generates ImageNet input_fn from a series of TFRecord files. 283 | 284 | The training data is assumed to be in TFRecord format with keys as specified 285 | in the dataset_parser below, sharded across 1024 files, named sequentially: 286 | 287 | train-00000-of-01024 288 | train-00001-of-01024 289 | ... 290 | train-01023-of-01024 291 | 292 | The validation data is in the same format but sharded in 128 files. 293 | 294 | The format of the data required is created by the script at: 295 | https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py 296 | """ 297 | 298 | def __init__(self, 299 | is_training, 300 | use_bfloat16, 301 | transpose_input, 302 | data_dir, 303 | image_size=224, 304 | num_parallel_calls=64, 305 | cache=False, 306 | num_label_classes=1000, 307 | include_background_label=False, 308 | augment_name=None, 309 | mixup_alpha=0.0, 310 | randaug_num_layers=None, 311 | randaug_magnitude=None, 312 | resize_method=None, 313 | holdout_shards=None): 314 | """Create an input from TFRecord files. 315 | 316 | Args: 317 | is_training: `bool` for whether the input is for training 318 | use_bfloat16: If True, use bfloat16 precision; else use float32. 319 | transpose_input: 'bool' for whether to use the double transpose trick 320 | data_dir: `str` for the directory of the training and validation data; 321 | if 'null' (the literal string 'null') or implicitly False 322 | then construct a null pipeline, consisting of empty images 323 | and blank labels. 324 | image_size: `int` for image size (both width and height). 325 | num_parallel_calls: concurrency level to use when reading data from disk. 326 | cache: if true, fill the dataset by repeating from its cache. 327 | num_label_classes: number of label classes. Default to 1000 for ImageNet. 328 | include_background_label: if true, label #0 is reserved for background. 329 | augment_name: `string` that is the name of the augmentation method 330 | to apply to the image. `autoaugment` if AutoAugment is to be used or 331 | `randaugment` if RandAugment is to be used. If the value is `None` no 332 | no augmentation method will be applied applied. See autoaugment.py 333 | for more details. 334 | mixup_alpha: float to control the strength of Mixup regularization, set 335 | to 0.0 to disable. 336 | randaug_num_layers: 'int', if RandAug is used, what should the number of 337 | layers be. See autoaugment.py for detailed description. 338 | randaug_magnitude: 'int', if RandAug is used, what should the magnitude 339 | be. See autoaugment.py for detailed description. 340 | resize_method: If None, use bicubic in default. 341 | holdout_shards: number of holdout training shards for validation. 342 | """ 343 | super(ImageNetInput, self).__init__( 344 | is_training=is_training, 345 | image_size=image_size, 346 | use_bfloat16=use_bfloat16, 347 | transpose_input=transpose_input, 348 | num_label_classes=num_label_classes, 349 | include_background_label=include_background_label, 350 | augment_name=augment_name, 351 | mixup_alpha=mixup_alpha, 352 | randaug_num_layers=randaug_num_layers, 353 | randaug_magnitude=randaug_magnitude) 354 | self.data_dir = data_dir 355 | if self.data_dir == 'null' or not self.data_dir: 356 | self.data_dir = None 357 | self.num_parallel_calls = num_parallel_calls 358 | self.cache = cache 359 | self.holdout_shards = holdout_shards 360 | 361 | def _get_null_input(self, data): 362 | """Returns a null image (all black pixels). 363 | 364 | Args: 365 | data: element of a dataset, ignored in this method, since it produces 366 | the same null image regardless of the element. 367 | 368 | Returns: 369 | a tensor representing a null image. 370 | """ 371 | del data # Unused since output is constant regardless of input 372 | return tf.zeros([self.image_size, self.image_size, 3], tf.bfloat16 373 | if self.use_bfloat16 else tf.float32) 374 | 375 | def dataset_parser(self, value): 376 | """See base class.""" 377 | if not self.data_dir: 378 | return value, tf.constant(0., tf.float32, (1000,)) 379 | return super(ImageNetInput, self).dataset_parser(value) 380 | 381 | def make_source_dataset(self, index, num_hosts): 382 | """See base class.""" 383 | if not self.data_dir: 384 | logging.info('Undefined data_dir implies null input') 385 | return tf.data.Dataset.range(1).repeat().map(self._get_null_input) 386 | 387 | if self.holdout_shards: 388 | if self.is_training: 389 | filenames = [ 390 | os.path.join(self.data_dir, 'train-%05d-of-01024' % i) 391 | for i in range(self.holdout_shards, 1024) 392 | ] 393 | else: 394 | filenames = [ 395 | os.path.join(self.data_dir, 'train-%05d-of-01024' % i) 396 | for i in range(0, self.holdout_shards) 397 | ] 398 | for f in filenames[:10]: 399 | logging.info('datafiles: %s', f) 400 | dataset = tf.data.Dataset.from_tensor_slices(filenames) 401 | else: 402 | file_pattern = os.path.join( 403 | self.data_dir, 'train-*' if self.is_training else 'validation-*') 404 | logging.info('datafiles: %s', file_pattern) 405 | dataset = tf.data.Dataset.list_files(file_pattern, shuffle=False) 406 | 407 | # For multi-host training, we want each hosts to always process the same 408 | # subset of files. Each host only sees a subset of the entire dataset, 409 | # allowing us to cache larger datasets in memory. 410 | dataset = dataset.shard(num_hosts, index) 411 | 412 | if self.is_training and not self.cache: 413 | dataset = dataset.repeat() 414 | 415 | def fetch_dataset(filename): 416 | buffer_size = 8 * 1024 * 1024 # 8 MiB per file 417 | dataset = tf.data.TFRecordDataset(filename, buffer_size=buffer_size) 418 | return dataset 419 | 420 | # Read the data from disk in parallel 421 | dataset = dataset.interleave( 422 | fetch_dataset, cycle_length=self.num_parallel_calls, 423 | num_parallel_calls=self.num_parallel_calls, deterministic=False) 424 | 425 | if self.cache: 426 | dataset = dataset.cache().shuffle(1024 * 16).repeat() 427 | else: 428 | dataset = dataset.shuffle(1024) 429 | return dataset 430 | 431 | 432 | # Defines a selection of data from a Cloud Bigtable. 433 | BigtableSelection = collections.namedtuple('BigtableSelection', [ 434 | 'project', 'instance', 'table', 'prefix', 'column_family', 435 | 'column_qualifier' 436 | ]) 437 | 438 | 439 | class ImageNetBigtableInput(ImageNetTFExampleInput): 440 | """Generates ImageNet input_fn from a Bigtable for training or evaluation. 441 | """ 442 | 443 | def __init__(self, 444 | is_training, 445 | use_bfloat16, 446 | transpose_input, 447 | selection, 448 | augment_name=None, 449 | num_label_classes=1000, 450 | include_background_label=False, 451 | mixup_alpha=0.0, 452 | randaug_num_layers=None, 453 | randaug_magnitude=None, 454 | resize_method=None): 455 | """Constructs an ImageNet input from a BigtableSelection. 456 | 457 | Args: 458 | is_training: `bool` for whether the input is for training 459 | use_bfloat16: If True, use bfloat16 precision; else use float32. 460 | transpose_input: 'bool' for whether to use the double transpose trick 461 | selection: a BigtableSelection specifying a part of a Bigtable. 462 | augment_name: `string` that is the name of the augmentation method 463 | to apply to the image. `autoaugment` if AutoAugment is to be used or 464 | `randaugment` if RandAugment is to be used. If the value is `None` no 465 | no augmentation method will be applied applied. See autoaugment.py 466 | for more details. 467 | num_label_classes: number of label classes. Default to 1000 for ImageNet. 468 | include_background_label: if true, label #0 is reserved for background. 469 | mixup_alpha: float to control the strength of Mixup regularization, set 470 | to 0.0 to disable. 471 | randaug_num_layers: 'int', if RandAug is used, what should the number of 472 | layers be. See autoaugment.py for detailed description. 473 | randaug_magnitude: 'int', if RandAug is used, what should the magnitude 474 | be. See autoaugment.py for detailed description.s 475 | resize_method: if None, use bicubic. 476 | """ 477 | super(ImageNetBigtableInput, self).__init__( 478 | is_training=is_training, 479 | use_bfloat16=use_bfloat16, 480 | transpose_input=transpose_input, 481 | num_label_classes=num_label_classes, 482 | include_background_label=include_background_label, 483 | augment_name=augment_name, 484 | mixup_alpha=mixup_alpha, 485 | randaug_num_layers=randaug_num_layers, 486 | randaug_magnitude=randaug_magnitude, 487 | resize_method=resize_method) 488 | self.selection = selection 489 | 490 | def make_source_dataset(self, index, num_hosts): 491 | """See base class.""" 492 | try: 493 | from tensorflow.contrib.cloud import BigtableClient # pylint: disable=g-import-not-at-top 494 | except ImportError as e: 495 | logging.exception('Bigtable is not supported in TensorFlow 2.x.') 496 | raise e 497 | 498 | data = self.selection 499 | client = BigtableClient(data.project, data.instance) 500 | table = client.table(data.table) 501 | ds = table.parallel_scan_prefix(data.prefix, 502 | columns=[(data.column_family, 503 | data.column_qualifier)]) 504 | # The Bigtable datasets will have the shape (row_key, data) 505 | ds_data = ds.map(lambda index, data: data) 506 | 507 | if self.is_training: 508 | ds_data = ds_data.repeat() 509 | 510 | return ds_data 511 | -------------------------------------------------------------------------------- /efficientnet_tf/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Model utilities.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import json 22 | import os 23 | import sys 24 | 25 | from absl import flags 26 | from absl import logging 27 | import numpy as np 28 | import tensorflow.compat.v1 as tf 29 | 30 | from . import lars_optimizer 31 | from tensorflow.python.tpu import tpu_function # pylint:disable=g-direct-tensorflow-import 32 | 33 | FLAGS = flags.FLAGS 34 | 35 | 36 | def build_learning_rate(initial_lr, 37 | global_step, 38 | steps_per_epoch=None, 39 | lr_decay_type='exponential', 40 | decay_factor=0.97, 41 | decay_epochs=2.4, 42 | total_steps=None, 43 | warmup_epochs=5): 44 | """Build learning rate.""" 45 | if lr_decay_type == 'exponential': 46 | assert steps_per_epoch is not None 47 | decay_steps = steps_per_epoch * decay_epochs 48 | lr = tf.train.exponential_decay( 49 | initial_lr, global_step, decay_steps, decay_factor, staircase=True) 50 | elif lr_decay_type == 'cosine': 51 | assert total_steps is not None 52 | lr = 0.5 * initial_lr * ( 53 | 1 + tf.cos(np.pi * tf.cast(global_step, tf.float32) / total_steps)) 54 | elif lr_decay_type == 'constant': 55 | lr = initial_lr 56 | elif lr_decay_type == 'poly': 57 | tf.logging.info('Using poly LR schedule') 58 | assert steps_per_epoch is not None 59 | assert total_steps is not None 60 | warmup_steps = int(steps_per_epoch * warmup_epochs) 61 | min_step = tf.constant(1, dtype=tf.int64) 62 | decay_steps = tf.maximum(min_step, tf.subtract(global_step, warmup_steps)) 63 | lr = tf.train.polynomial_decay( 64 | initial_lr, 65 | decay_steps, 66 | total_steps - warmup_steps + 1, 67 | end_learning_rate=0.1, 68 | power=2.0) 69 | else: 70 | assert False, 'Unknown lr_decay_type : %s' % lr_decay_type 71 | 72 | if warmup_epochs: 73 | logging.info('Learning rate warmup_epochs: %d', warmup_epochs) 74 | warmup_steps = int(warmup_epochs * steps_per_epoch) 75 | warmup_lr = ( 76 | initial_lr * tf.cast(global_step, tf.float32) / tf.cast( 77 | warmup_steps, tf.float32)) 78 | lr = tf.cond(global_step < warmup_steps, lambda: warmup_lr, lambda: lr) 79 | 80 | return lr 81 | 82 | 83 | def build_optimizer(learning_rate, 84 | optimizer_name='rmsprop', 85 | decay=0.9, 86 | epsilon=0.001, 87 | momentum=0.9, 88 | lars_weight_decay=None, 89 | lars_epsilon=None): 90 | """Build optimizer.""" 91 | if optimizer_name == 'sgd': 92 | logging.info('Using SGD optimizer') 93 | optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate) 94 | elif optimizer_name == 'momentum': 95 | logging.info('Using Momentum optimizer') 96 | optimizer = tf.train.MomentumOptimizer( 97 | learning_rate=learning_rate, momentum=momentum) 98 | elif optimizer_name == 'rmsprop': 99 | logging.info('Using RMSProp optimizer') 100 | optimizer = tf.train.RMSPropOptimizer(learning_rate, decay, momentum, 101 | epsilon) 102 | elif optimizer_name == 'lars': 103 | logging.info('Using LARS optimizer') 104 | assert lars_weight_decay is not None, 'LARS weight decay is None.' 105 | assert lars_epsilon is not None, 'LARS epsilon is None.' 106 | optimizer = lars_optimizer.LARSOptimizer( 107 | learning_rate, 108 | momentum=momentum, 109 | weight_decay=lars_weight_decay, 110 | skip_list=['batch_normalization', 'bias', 'beta', 'gamma'], 111 | epsilon=lars_epsilon) 112 | else: 113 | logging.fatal('Unknown optimizer: %s', optimizer_name) 114 | 115 | return optimizer 116 | 117 | 118 | class TpuBatchNormalization(tf.layers.BatchNormalization): 119 | # class TpuBatchNormalization(tf.layers.BatchNormalization): 120 | """Cross replica batch normalization.""" 121 | 122 | def __init__(self, fused=False, **kwargs): 123 | if fused in (True, None): 124 | raise ValueError('TpuBatchNormalization does not support fused=True.') 125 | super(TpuBatchNormalization, self).__init__(fused=fused, **kwargs) 126 | 127 | def _cross_replica_average(self, t, num_shards_per_group): 128 | """Calculates the average value of input tensor across TPU replicas.""" 129 | num_shards = tpu_function.get_tpu_context().number_of_shards 130 | group_assignment = None 131 | if num_shards_per_group > 1: 132 | if num_shards % num_shards_per_group != 0: 133 | raise ValueError('num_shards: %d mod shards_per_group: %d, should be 0' 134 | % (num_shards, num_shards_per_group)) 135 | num_groups = num_shards // num_shards_per_group 136 | group_assignment = [[ 137 | x for x in range(num_shards) if x // num_shards_per_group == y 138 | ] for y in range(num_groups)] 139 | return tf.tpu.cross_replica_sum(t, group_assignment) / tf.cast( 140 | num_shards_per_group, t.dtype) 141 | 142 | def _moments(self, inputs, reduction_axes, keep_dims): 143 | """Compute the mean and variance: it overrides the original _moments.""" 144 | shard_mean, shard_variance = super(TpuBatchNormalization, self)._moments( 145 | inputs, reduction_axes, keep_dims=keep_dims) 146 | 147 | num_shards = tpu_function.get_tpu_context().number_of_shards or 1 148 | if num_shards <= 8: # Skip cross_replica for 2x2 or smaller slices. 149 | num_shards_per_group = 1 150 | else: 151 | num_shards_per_group = max(8, num_shards // 8) 152 | logging.info('TpuBatchNormalization with num_shards_per_group %s', 153 | num_shards_per_group) 154 | if num_shards_per_group > 1: 155 | # Compute variance using: Var[X]= E[X^2] - E[X]^2. 156 | shard_square_of_mean = tf.math.square(shard_mean) 157 | shard_mean_of_square = shard_variance + shard_square_of_mean 158 | group_mean = self._cross_replica_average( 159 | shard_mean, num_shards_per_group) 160 | group_mean_of_square = self._cross_replica_average( 161 | shard_mean_of_square, num_shards_per_group) 162 | group_variance = group_mean_of_square - tf.math.square(group_mean) 163 | return (group_mean, group_variance) 164 | else: 165 | return (shard_mean, shard_variance) 166 | 167 | 168 | class BatchNormalization(tf.layers.BatchNormalization): 169 | """Fixed default name of BatchNormalization to match TpuBatchNormalization.""" 170 | 171 | def __init__(self, name='tpu_batch_normalization', **kwargs): 172 | super(BatchNormalization, self).__init__(name=name, **kwargs) 173 | 174 | 175 | def train_batch_norm(**kwargs): 176 | if 'optimizer' in FLAGS and FLAGS.optimizer == 'lars': 177 | return DistributedBatchNormalization(**kwargs) 178 | return TpuBatchNormalization(**kwargs) 179 | 180 | 181 | def eval_batch_norm(**kwargs): 182 | if 'optimizer' in FLAGS and FLAGS.optimizer == 'lars': 183 | return DistributedBatchNormalization(**kwargs) 184 | return BatchNormalization(**kwargs) 185 | 186 | 187 | class DistributedBatchNormalization: 188 | """Distributed batch normalization used in https://arxiv.org/abs/2011.00071.""" 189 | 190 | def __init__(self, axis, momentum, epsilon): 191 | self.axis = axis 192 | self.momentum = momentum 193 | self.epsilon = epsilon 194 | 195 | def __call__(self, x, training, distname='batch_normalization'): 196 | shape = [x.shape[-1]] 197 | with tf.variable_scope('batch_normalization'): 198 | ones = tf.initializers.ones() 199 | zeros = tf.initializers.zeros() 200 | gamma = tf.get_variable( 201 | 'gamma', shape, initializer=ones, trainable=True, use_resource=True) 202 | beta = tf.get_variable( 203 | 'beta', shape, initializer=zeros, trainable=True, use_resource=True) 204 | moving_mean = tf.get_variable( 205 | 'moving_mean', 206 | shape, 207 | initializer=zeros, 208 | trainable=False, 209 | use_resource=True) 210 | moving_variance = tf.get_variable( 211 | 'moving_variance', 212 | shape, 213 | initializer=ones, 214 | trainable=False, 215 | use_resource=True) 216 | num_replicas = FLAGS.num_replicas 217 | 218 | x = tf.cast(x, tf.float32) 219 | if training: 220 | if num_replicas <= 8: 221 | group_assign = None 222 | group_shards = tf.cast(num_replicas, tf.float32) 223 | else: 224 | 225 | group_shards = max( 226 | 1, 227 | int(FLAGS.batch_norm_batch_size / 228 | (FLAGS.train_batch_size / num_replicas))) 229 | group_assign = np.arange(num_replicas, dtype=np.int32) 230 | group_assign = group_assign.reshape([-1, group_shards]) 231 | group_assign = group_assign.tolist() 232 | group_shards = tf.cast(group_shards, tf.float32) 233 | 234 | mean = tf.reduce_mean(x, [0, 1, 2]) 235 | mean = tf.tpu.cross_replica_sum(mean, group_assign) / group_shards 236 | 237 | # Var[x] = E[x^2] - E[x]^2 238 | mean_sq = tf.reduce_mean(tf.math.square(x), [0, 1, 2]) 239 | mean_sq = tf.tpu.cross_replica_sum(mean_sq, group_assign) / group_shards 240 | variance = mean_sq - tf.math.square(mean) 241 | 242 | decay = tf.cast(1. - self.momentum, tf.float32) 243 | 244 | def u(moving, normal, name): 245 | num_replicas_fp = tf.cast(num_replicas, tf.float32) 246 | normal = tf.tpu.cross_replica_sum(normal) / num_replicas_fp 247 | diff = decay * (moving - normal) 248 | return tf.assign_sub(moving, diff, use_locking=True, name=name) 249 | 250 | tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, 251 | u(moving_mean, mean, name='moving_mean')) 252 | tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, 253 | u(moving_variance, variance, name='moving_variance')) 254 | 255 | x = tf.nn.batch_normalization( 256 | x, 257 | mean=mean, 258 | variance=variance, 259 | offset=beta, 260 | scale=gamma, 261 | variance_epsilon=self.epsilon) 262 | else: 263 | 264 | x, _, _ = tf.nn.fused_batch_norm( 265 | x, 266 | scale=gamma, 267 | offset=beta, 268 | mean=moving_mean, 269 | variance=moving_variance, 270 | epsilon=self.epsilon, 271 | is_training=False) 272 | 273 | return x 274 | 275 | 276 | def drop_connect(inputs, is_training, survival_prob): 277 | """Drop the entire conv with given survival probability.""" 278 | # "Deep Networks with Stochastic Depth", https://arxiv.org/pdf/1603.09382.pdf 279 | if not is_training: 280 | return inputs 281 | 282 | # Compute tensor. 283 | batch_size = tf.shape(inputs)[0] 284 | random_tensor = survival_prob 285 | random_tensor += tf.random_uniform([batch_size, 1, 1, 1], dtype=inputs.dtype) 286 | binary_tensor = tf.floor(random_tensor) 287 | # Unlike conventional way that multiply survival_prob at test time, here we 288 | # divide survival_prob at training time, such that no addition compute is 289 | # needed at test time. 290 | output = tf.div(inputs, survival_prob) * binary_tensor 291 | return output 292 | 293 | 294 | def archive_ckpt(ckpt_eval, ckpt_objective, ckpt_path): 295 | """Archive a checkpoint if the metric is better.""" 296 | ckpt_dir, ckpt_name = os.path.split(ckpt_path) 297 | 298 | saved_objective_path = os.path.join(ckpt_dir, 'best_objective.txt') 299 | saved_objective = float('-inf') 300 | if tf.gfile.Exists(saved_objective_path): 301 | with tf.gfile.GFile(saved_objective_path, 'r') as f: 302 | saved_objective = float(f.read()) 303 | if saved_objective > ckpt_objective: 304 | logging.info('Ckpt %s is worse than %s', ckpt_objective, saved_objective) 305 | return False 306 | 307 | filenames = tf.gfile.Glob(ckpt_path + '.*') 308 | if filenames is None: 309 | logging.info('No files to copy for checkpoint %s', ckpt_path) 310 | return False 311 | 312 | # Clear the old folder. 313 | dst_dir = os.path.join(ckpt_dir, 'archive') 314 | if tf.gfile.Exists(dst_dir): 315 | tf.gfile.DeleteRecursively(dst_dir) 316 | tf.gfile.MakeDirs(dst_dir) 317 | 318 | # Write checkpoints. 319 | for f in filenames: 320 | dest = os.path.join(dst_dir, os.path.basename(f)) 321 | tf.gfile.Copy(f, dest, overwrite=True) 322 | ckpt_state = tf.train.generate_checkpoint_state_proto( 323 | dst_dir, 324 | model_checkpoint_path=ckpt_name, 325 | all_model_checkpoint_paths=[ckpt_name]) 326 | with tf.gfile.GFile(os.path.join(dst_dir, 'checkpoint'), 'w') as f: 327 | f.write(str(ckpt_state)) 328 | with tf.gfile.GFile(os.path.join(dst_dir, 'best_eval.txt'), 'w') as f: 329 | f.write('%s' % ckpt_eval) 330 | 331 | # Update the best objective. 332 | with tf.gfile.GFile(saved_objective_path, 'w') as f: 333 | f.write('%f' % ckpt_objective) 334 | 335 | logging.info('Copying checkpoint %s to %s', ckpt_path, dst_dir) 336 | return True 337 | 338 | 339 | def get_ema_vars(): 340 | """Get all exponential moving average (ema) variables.""" 341 | ema_vars = tf.trainable_variables() + tf.get_collection('moving_vars') 342 | for v in tf.global_variables(): 343 | # We maintain mva for batch norm moving mean and variance as well. 344 | if 'moving_mean' in v.name or 'moving_variance' in v.name: 345 | ema_vars.append(v) 346 | return list(set(ema_vars)) 347 | 348 | 349 | class DepthwiseConv2D(tf.keras.layers.DepthwiseConv2D, tf.layers.Layer): 350 | """Wrap keras DepthwiseConv2D to tf.layers.""" 351 | 352 | pass 353 | 354 | 355 | class Conv2D(tf.layers.Conv2D): 356 | """Wrapper for Conv2D with specialization for fast inference.""" 357 | 358 | def _bias_activation(self, outputs): 359 | if self.use_bias: 360 | outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW') 361 | if self.activation is not None: 362 | return self.activation(outputs) 363 | return outputs 364 | 365 | def _can_run_fast_1x1(self, inputs): 366 | batch_size = inputs.shape.as_list()[0] 367 | return (self.data_format == 'channels_first' and 368 | batch_size == 1 and 369 | self.kernel_size == (1, 1)) 370 | 371 | def _call_fast_1x1(self, inputs): 372 | # Compute the 1x1 convolution as a matmul. 373 | inputs_shape = tf.shape(inputs) 374 | flat_inputs = tf.reshape(inputs, [inputs_shape[1], -1]) 375 | flat_outputs = tf.matmul( 376 | tf.squeeze(self.kernel), 377 | flat_inputs, 378 | transpose_a=True) 379 | outputs_shape = tf.concat([[1, self.filters], inputs_shape[2:]], axis=0) 380 | outputs = tf.reshape(flat_outputs, outputs_shape) 381 | 382 | # Handle the bias and activation function. 383 | return self._bias_activation(outputs) 384 | 385 | def call(self, inputs): 386 | if self._can_run_fast_1x1(inputs): 387 | return self._call_fast_1x1(inputs) 388 | return super(Conv2D, self).call(inputs) 389 | 390 | 391 | class EvalCkptDriver(object): 392 | """A driver for running eval inference. 393 | 394 | Attributes: 395 | model_name: str. Model name to eval. 396 | batch_size: int. Eval batch size. 397 | image_size: int. Input image size, determined by model name. 398 | num_classes: int. Number of classes, default to 1000 for ImageNet. 399 | include_background_label: whether to include extra background label. 400 | advprop_preprocessing: whether to use advprop preprocessing. 401 | """ 402 | 403 | def __init__(self, 404 | model_name, 405 | batch_size=1, 406 | image_size=224, 407 | num_classes=1000, 408 | include_background_label=False, 409 | advprop_preprocessing=False): 410 | """Initialize internal variables.""" 411 | self.model_name = model_name 412 | self.batch_size = batch_size 413 | self.num_classes = num_classes 414 | self.include_background_label = include_background_label 415 | self.image_size = image_size 416 | self.advprop_preprocessing = advprop_preprocessing 417 | 418 | def restore_model(self, sess, ckpt_dir, enable_ema=True, export_ckpt=None): 419 | """Restore variables from checkpoint dir.""" 420 | sess.run(tf.global_variables_initializer()) 421 | checkpoint = tf.train.latest_checkpoint(ckpt_dir) 422 | if enable_ema: 423 | ema = tf.train.ExponentialMovingAverage(decay=0.0) 424 | ema_vars = get_ema_vars() 425 | var_dict = ema.variables_to_restore(ema_vars) 426 | ema_assign_op = ema.apply(ema_vars) 427 | else: 428 | var_dict = get_ema_vars() 429 | ema_assign_op = None 430 | 431 | tf.train.get_or_create_global_step() 432 | sess.run(tf.global_variables_initializer()) 433 | saver = tf.train.Saver(var_dict, max_to_keep=1) 434 | saver.restore(sess, checkpoint) 435 | 436 | if export_ckpt: 437 | if ema_assign_op is not None: 438 | sess.run(ema_assign_op) 439 | saver = tf.train.Saver(max_to_keep=1, save_relative_paths=True) 440 | saver.save(sess, export_ckpt) 441 | 442 | def build_model(self, features, is_training): 443 | """Build model with input features.""" 444 | del features, is_training 445 | raise ValueError('Must be implemented by subclasses.') 446 | 447 | def get_preprocess_fn(self): 448 | raise ValueError('Must be implemented by subclsses.') 449 | 450 | def build_dataset(self, filenames, labels, is_training): 451 | """Build input dataset.""" 452 | batch_drop_remainder = False 453 | if 'condconv' in self.model_name and not is_training: 454 | # CondConv layers can only be called with known batch dimension. Thus, we 455 | # must drop all remaining examples that do not make up one full batch. 456 | # To ensure all examples are evaluated, use a batch size that evenly 457 | # divides the number of files. 458 | batch_drop_remainder = True 459 | num_files = len(filenames) 460 | if num_files % self.batch_size != 0: 461 | tf.logging.warn('Remaining examples in last batch are not being ' 462 | 'evaluated.') 463 | filenames = tf.constant(filenames) 464 | labels = tf.constant(labels) 465 | dataset = tf.data.Dataset.from_tensor_slices((filenames, labels)) 466 | 467 | def _parse_function(filename, label): 468 | image_string = tf.read_file(filename) 469 | preprocess_fn = self.get_preprocess_fn() 470 | image_decoded = preprocess_fn( 471 | image_string, is_training, image_size=self.image_size) 472 | image = tf.cast(image_decoded, tf.float32) 473 | return image, label 474 | 475 | dataset = dataset.map(_parse_function) 476 | dataset = dataset.batch(self.batch_size, 477 | drop_remainder=batch_drop_remainder) 478 | 479 | iterator = dataset.make_one_shot_iterator() 480 | images, labels = iterator.get_next() 481 | return images, labels 482 | 483 | def run_inference(self, 484 | ckpt_dir, 485 | image_files, 486 | labels, 487 | enable_ema=True, 488 | export_ckpt=None): 489 | """Build and run inference on the target images and labels.""" 490 | label_offset = 1 if self.include_background_label else 0 491 | with tf.Graph().as_default(), tf.Session() as sess: 492 | images, labels = self.build_dataset(image_files, labels, False) 493 | probs = self.build_model(images, is_training=False) 494 | if isinstance(probs, tuple): 495 | probs = probs[0] 496 | 497 | self.restore_model(sess, ckpt_dir, enable_ema, export_ckpt) 498 | 499 | prediction_idx = [] 500 | prediction_prob = [] 501 | for _ in range(len(image_files) // self.batch_size): 502 | out_probs = sess.run(probs) 503 | idx = np.argsort(out_probs)[::-1] 504 | prediction_idx.append(idx[:5] - label_offset) 505 | prediction_prob.append([out_probs[pid] for pid in idx[:5]]) 506 | 507 | # Return the top 5 predictions (idx and prob) for each image. 508 | return prediction_idx, prediction_prob 509 | 510 | def eval_example_images(self, 511 | ckpt_dir, 512 | image_files, 513 | labels_map_file, 514 | enable_ema=True, 515 | export_ckpt=None): 516 | """Eval a list of example images. 517 | 518 | Args: 519 | ckpt_dir: str. Checkpoint directory path. 520 | image_files: List[str]. A list of image file paths. 521 | labels_map_file: str. The labels map file path. 522 | enable_ema: enable expotential moving average. 523 | export_ckpt: export ckpt folder. 524 | 525 | Returns: 526 | A tuple (pred_idx, and pred_prob), where pred_idx is the top 5 prediction 527 | index and pred_prob is the top 5 prediction probability. 528 | """ 529 | classes = json.loads(tf.gfile.Open(labels_map_file).read()) 530 | pred_idx, pred_prob = self.run_inference( 531 | ckpt_dir, image_files, [0] * len(image_files), enable_ema, export_ckpt) 532 | for i in range(len(image_files)): 533 | print('predicted class for image {}: '.format(image_files[i])) 534 | for j, idx in enumerate(pred_idx[i]): 535 | print(' -> top_{} ({:4.2f}%): {} '.format(j, pred_prob[i][j] * 100, 536 | classes[str(idx)])) 537 | return pred_idx, pred_prob 538 | 539 | def eval_imagenet(self, ckpt_dir, imagenet_eval_glob, 540 | imagenet_eval_label, num_images, enable_ema, export_ckpt): 541 | """Eval ImageNet images and report top1/top5 accuracy. 542 | 543 | Args: 544 | ckpt_dir: str. Checkpoint directory path. 545 | imagenet_eval_glob: str. File path glob for all eval images. 546 | imagenet_eval_label: str. File path for eval label. 547 | num_images: int. Number of images to eval: -1 means eval the whole 548 | dataset. 549 | enable_ema: enable expotential moving average. 550 | export_ckpt: export checkpoint folder. 551 | 552 | Returns: 553 | A tuple (top1, top5) for top1 and top5 accuracy. 554 | """ 555 | imagenet_val_labels = [int(i) for i in tf.gfile.GFile(imagenet_eval_label)] 556 | imagenet_filenames = sorted(tf.gfile.Glob(imagenet_eval_glob)) 557 | if num_images < 0: 558 | num_images = len(imagenet_filenames) 559 | image_files = imagenet_filenames[:num_images] 560 | labels = imagenet_val_labels[:num_images] 561 | 562 | pred_idx, _ = self.run_inference( 563 | ckpt_dir, image_files, labels, enable_ema, export_ckpt) 564 | top1_cnt, top5_cnt = 0.0, 0.0 565 | for i, label in enumerate(labels): 566 | top1_cnt += label in pred_idx[i][:1] 567 | top5_cnt += label in pred_idx[i][:5] 568 | if i % 100 == 0: 569 | print('Step {}: top1_acc = {:4.2f}% top5_acc = {:4.2f}%'.format( 570 | i, 100 * top1_cnt / (i + 1), 100 * top5_cnt / (i + 1))) 571 | sys.stdout.flush() 572 | top1, top5 = 100 * top1_cnt / num_images, 100 * top5_cnt / num_images 573 | print('Final: top1_acc = {:4.2f}% top5_acc = {:4.2f}%'.format(top1, top5)) 574 | return top1, top5 575 | --------------------------------------------------------------------------------