├── .gitignore ├── LICENSE ├── README.md ├── anynet_cfg.py ├── configuration.py ├── dataset └── README.md ├── evaluate.py ├── models ├── RegNet │ ├── __init__.py │ ├── anynet.py │ ├── blocks.py │ └── regnet.py ├── __init__.py ├── densenet.py ├── efficientnet.py ├── group_convolution.py ├── inception_modules.py ├── inception_resnet_v1.py ├── inception_resnet_v2.py ├── inception_v4.py ├── mobilenet_v1.py ├── mobilenet_v2.py ├── mobilenet_v3_block.py ├── mobilenet_v3_large.py ├── mobilenet_v3_small.py ├── residual_block.py ├── resnet.py ├── resnext.py ├── resnext_block.py ├── se_resnet.py ├── se_resnext.py ├── shufflenet_v2.py └── squeezenet.py ├── original_dataset └── README.md ├── parse_tfrecord.py ├── predict.py ├── prepare_data.py ├── saved_model └── README.md ├── show_model_list.py ├── split_dataset.py ├── test_single_image.py ├── to_tfrecord.py └── train.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | __pycache__ 3 | /dataset/* 4 | !/dataset/README.md 5 | /original_dataset/* 6 | !/original_dataset/README.md 7 | /saved_model/* 8 | !/saved_model/README.md 9 | models.txt 10 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 calmisential 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Basic_CNNs_TensorFlow2 2 | A tensorflow2 implementation of some basic CNNs. 3 | 4 | ## Networks included: 5 | + MobileNet_V1 6 | + MobileNet_V2 7 | + [MobileNet_V3](https://github.com/calmisential/MobileNetV3_TensorFlow2) 8 | + [EfficientNet](https://github.com/calmisential/EfficientNet_TensorFlow2) 9 | + [ResNeXt](https://github.com/calmisential/ResNeXt_TensorFlow2) 10 | + [InceptionV4, InceptionResNetV1, InceptionResNetV2](https://github.com/calmisential/InceptionV4_TensorFlow2) 11 | + SE_ResNet_50, SE_ResNet_101, SE_ResNet_152, SE_ResNeXt_50, SE_ResNeXt_101 12 | + SqueezeNet 13 | + [DenseNet](https://github.com/calmisential/DenseNet_TensorFlow2) 14 | + ShuffleNetV2 15 | + [ResNet](https://github.com/calmisential/TensorFlow2.0_ResNet) 16 | + RegNet 17 | 18 | ## Other networks 19 | For AlexNet and VGG, see : https://github.com/calmisential/TensorFlow2.0_Image_Classification
20 | For InceptionV3, see : https://github.com/calmisential/TensorFlow2.0_InceptionV3
21 | For ResNet, see : https://github.com/calmisential/TensorFlow2.0_ResNet 22 | 23 | ## Train 24 | 1. Requirements: 25 | + Python >= 3.9 26 | + Tensorflow >= 2.7.0 27 | + tensorflow-addons >= 0.15.0 28 | 2. To train the network on your own dataset, you can put the dataset under the folder **original dataset**, and the directory should look like this: 29 | ``` 30 | |——original dataset 31 | |——class_name_0 32 | |——class_name_1 33 | |——class_name_2 34 | |——class_name_3 35 | ``` 36 | 3. Run the script **split_dataset.py** to split the raw dataset into train set, valid set and test set. The dataset directory will be like this: 37 | ``` 38 | |——dataset 39 | |——train 40 | |——class_name_1 41 | |——class_name_2 42 | ...... 43 | |——class_name_n 44 | |——valid 45 | |——class_name_1 46 | |——class_name_2 47 | ...... 48 | |——class_name_n 49 | |—-test 50 | |——class_name_1 51 | |——class_name_2 52 | ...... 53 | |——class_name_n 54 | ``` 55 | 4. Run **to_tfrecord.py** to generate tfrecord files. 56 | 5. Change the corresponding parameters in **config.py**. 57 | 6. Run **show_model_list.py** to get the index of model. 58 | 7. Run **python train.py --idx [index]** to start training.
59 | If you want to train the *EfficientNet*, you should change the IMAGE_HEIGHT and IMAGE_WIDTH before training. 60 | - b0 = (224, 224) 61 | - b1 = (240, 240) 62 | - b2 = (260, 260) 63 | - b3 = (300, 300) 64 | - b4 = (380, 380) 65 | - b5 = (456, 456) 66 | - b6 = (528, 528) 67 | - b7 = (600, 600) 68 | ## Evaluate 69 | Run **python evaluate.py --idx [index]** to evaluate the model's performance on the test dataset. 70 | 71 | ## Different input image sizes for different neural networks 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | 139 | 140 | 141 | 142 | 143 | 144 | 145 | 146 | 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | 160 | 161 | 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 |
TypeNeural NetworkInput Image Size (height * width)
MobileNetMobileNet_V1(224 * 224)
MobileNet_V2(224 * 224)
MobileNet_V3(224 * 224)
EfficientNetEfficientNet(B0~B7)/
ResNeXtResNeXt50(224 * 224)
ResNeXt101(224 * 224)
SEResNeXtSEResNeXt50(224 * 224)
SEResNeXt101(224 * 224)
InceptionInceptionV4(299 * 299)
Inception_ResNet_V1(299 * 299)
Inception_ResNet_V2(299 * 299)
SE_ResNetSE_ResNet_50(224 * 224)
SE_ResNet_101(224 * 224)
SE_ResNet_152(224 * 224)
SqueezeNetSqueezeNet(224 * 224)
DenseNetDenseNet_121(224 * 224)
DenseNet_169(224 * 224)
DenseNet_201(224 * 224)
DenseNet_269(224 * 224)
ShuffleNetV2ShuffleNetV2(224 * 224)
ResNetResNet_18(224 * 224)
ResNet_34(224 * 224)
ResNet_50(224 * 224)
ResNet_101(224 * 224)
ResNet_152(224 * 224)
189 | 190 | ## References 191 | 1. MobileNet_V1: [Efficient Convolutional Neural Networks for Mobile Vision Applications](https://arxiv.org/abs/1704.04861) 192 | 2. MobileNet_V2: [Inverted Residuals and Linear Bottlenecks](https://arxiv.org/abs/1801.04381) 193 | 3. MobileNet_V3: [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244) 194 | 4. EfficientNet: [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) 195 | 5. The official code of EfficientNet: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet 196 | 6. ResNeXt: [Aggregated Residual Transformations for Deep Neural Networks](https://arxiv.org/abs/1611.05431) 197 | 7. Inception_V4/Inception_ResNet_V1/Inception_ResNet_V2: [Inception-v4, Inception-ResNet and the Impact of Residual Connectionson Learning](https://arxiv.org/abs/1602.07261) 198 | 8. The official implementation of Inception_V4: https://github.com/tensorflow/models/blob/master/research/slim/nets/inception_v4.py 199 | 9. The official implementation of Inception_ResNet_V2: https://github.com/tensorflow/models/blob/master/research/slim/nets/inception_resnet_v2.py 200 | 10. SENet: [Squeeze-and-Excitation Networks](https://arxiv.org/abs/1709.01507) 201 | 11. SqueezeNet: [SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size](https://arxiv.org/abs/1602.07360) 202 | 12. DenseNet: [Densely Connected Convolutional Networks](https://arxiv.org/abs/1608.06993) 203 | 13. https://zhuanlan.zhihu.com/p/37189203 204 | 14. ShuffleNetV2: [ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design 205 | ](https://arxiv.org/abs/1807.11164) 206 | 15. https://zhuanlan.zhihu.com/p/48261931 207 | 16. ResNet: [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385) 208 | 17. RegNet: [Designing Network Design Spaces](https://arxiv.org/abs/2003.13678) -------------------------------------------------------------------------------- /anynet_cfg.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | class AnyNetCfg: 4 | stem_type = "simple_stem_in" 5 | stem_w = 32 6 | block_type = "res_bottleneck_block" 7 | depths = [] 8 | widths = [] 9 | strides = [] 10 | bot_muls = [] 11 | group_ws = [] 12 | se_on = True 13 | se_r = 0.25 14 | 15 | 16 | class RegNetCfg: 17 | stem_type = "simple_stem_in" 18 | stem_w = 32 19 | block_type = "res_bottleneck_block" 20 | stride = 2 21 | se_on = True 22 | se_r = 0.25 23 | depth = 10 24 | w0 = 32 25 | wa = 5.0 26 | wm = 2.5 27 | group_w = 16 28 | bot_mul = 1.0 29 | -------------------------------------------------------------------------------- /configuration.py: -------------------------------------------------------------------------------- 1 | 2 | DEVICE = "gpu" # cpu or gpu 3 | 4 | # some training parameters 5 | EPOCHS = 50 6 | BATCH_SIZE = 8 7 | NUM_CLASSES = 5 8 | IMAGE_HEIGHT = 224 9 | IMAGE_WIDTH = 224 10 | CHANNELS = 3 11 | 12 | save_model_dir = "saved_model/" 13 | save_every_n_epoch = 10 14 | test_image_dir = "" 15 | 16 | dataset_dir = "dataset/" 17 | train_dir = dataset_dir + "train" 18 | valid_dir = dataset_dir + "valid" 19 | test_dir = dataset_dir + "test" 20 | train_tfrecord = dataset_dir + "train.tfrecord" 21 | valid_tfrecord = dataset_dir + "valid.tfrecord" 22 | test_tfrecord = dataset_dir + "test.tfrecord" 23 | # VALID_SET_RATIO = 1 - TRAIN_SET_RATIO - TEST_SET_RATIO 24 | TRAIN_SET_RATIO = 0.6 25 | TEST_SET_RATIO = 0.2 26 | 27 | 28 | -------------------------------------------------------------------------------- /dataset/README.md: -------------------------------------------------------------------------------- 1 | The dataset includes train set, valid set and test set. -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import tensorflow as tf 4 | from configuration import save_model_dir 5 | from prepare_data import generate_datasets 6 | from train import process_features 7 | from models import get_model 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument("--idx", default=0, type=int) 11 | 12 | if __name__ == '__main__': 13 | gpus = tf.config.list_physical_devices('GPU') 14 | if gpus: 15 | try: 16 | for gpu in gpus: 17 | tf.config.experimental.set_memory_growth(gpu, True) 18 | logical_gpus = tf.config.list_logical_devices('GPU') 19 | print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") 20 | except RuntimeError as e: 21 | print(e) 22 | 23 | args = parser.parse_args() 24 | 25 | # get the original_dataset 26 | train_dataset, valid_dataset, test_dataset, train_count, valid_count, test_count = generate_datasets() 27 | # load the model 28 | model = get_model(args.idx) 29 | model.load_weights(filepath=save_model_dir) 30 | # model = tf.saved_model.load(save_model_dir) 31 | 32 | # Get the accuracy on the test set 33 | loss_object = tf.keras.metrics.SparseCategoricalCrossentropy() 34 | test_loss = tf.keras.metrics.Mean() 35 | test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy() 36 | 37 | # @tf.function 38 | def test_step(images, labels): 39 | predictions = model(images, training=False) 40 | t_loss = loss_object(labels, predictions) 41 | test_loss(t_loss) 42 | test_accuracy(labels, predictions) 43 | 44 | for features in test_dataset: 45 | test_images, test_labels = process_features(features, data_augmentation=False) 46 | test_step(test_images, test_labels) 47 | print("loss: {:.5f}, test accuracy: {:.5f}".format(test_loss.result(), 48 | test_accuracy.result())) 49 | 50 | print("The accuracy on test set is: {:.3f}%".format(test_accuracy.result()*100)) -------------------------------------------------------------------------------- /models/RegNet/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/calmiLovesAI/Basic_CNNs_TensorFlow2/f063c84451f12e904f9c91c51278be52afccb0c2/models/RegNet/__init__.py -------------------------------------------------------------------------------- /models/RegNet/anynet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import tensorflow_addons as tfa 3 | 4 | from models.RegNet.blocks import SE 5 | from anynet_cfg import AnyNetCfg 6 | from configuration import NUM_CLASSES 7 | from models.group_convolution import get_group_conv 8 | 9 | 10 | def get_stem_fun(stem_type): 11 | stem_funs = { 12 | "res_stem_cifar": ResStemCifar, 13 | "res_stem_in": ResStem, 14 | "simple_stem_in": SimpleStem, 15 | } 16 | err_str = "Stem type '{}' not supported" 17 | assert stem_type in stem_funs.keys(), err_str.format(stem_type) 18 | return stem_funs[stem_type] 19 | 20 | 21 | def get_block_fun(block_type): 22 | block_funs = { 23 | "vanilla_block": VanillaBlock, 24 | "res_basic_block": ResBasicBlock, 25 | "res_bottleneck_block": ResBottleneckBlock, 26 | } 27 | err_str = "Block type '{}' not supported" 28 | assert block_type in block_funs.keys(), err_str.format(block_type) 29 | return block_funs[block_type] 30 | 31 | 32 | class AnyHead(tf.keras.layers.Layer): 33 | def __init__(self, num_classes): 34 | super(AnyHead, self).__init__() 35 | self.avg_pool = tfa.layers.AdaptiveAveragePooling2D(output_size=(1, 1)) 36 | self.flatten = tf.keras.layers.Flatten() 37 | self.fc = tf.keras.layers.Dense(units=num_classes, use_bias=True) 38 | 39 | def call(self, inputs, **kwargs): 40 | x = self.avg_pool(inputs) 41 | x = self.flatten(x) 42 | x = self.fc(x) 43 | return x 44 | 45 | 46 | class VanillaBlock(tf.keras.layers.Layer): 47 | def __init__(self, w_out, stride): 48 | super(VanillaBlock, self).__init__() 49 | self.a = tf.keras.layers.Conv2D(filters=w_out, kernel_size=3, strides=stride, padding="same") 50 | self.a_bn = tf.keras.layers.BatchNormalization() 51 | self.b = tf.keras.layers.Conv2D(filters=w_out, kernel_size=3, padding="same") 52 | self.b_bn = tf.keras.layers.BatchNormalization() 53 | 54 | def call(self, inputs, training=None, **kwargs): 55 | x = self.a(inputs) 56 | x = self.a_bn(x, training=training) 57 | x = tf.keras.activations.relu(x) 58 | x = self.b(x) 59 | x = self.b_bn(x, training=training) 60 | x = tf.keras.activations.relu(x) 61 | return x 62 | 63 | 64 | class BasicTransform(tf.keras.layers.Layer): 65 | def __init__(self, w_out, stride): 66 | super(BasicTransform, self).__init__() 67 | self.a = tf.keras.layers.Conv2D(filters=w_out, kernel_size=3, strides=stride, padding="same") 68 | self.a_bn = tf.keras.layers.BatchNormalization() 69 | self.b = tf.keras.layers.Conv2D(filters=w_out, kernel_size=3, padding="same") 70 | self.b_bn = tf.keras.layers.BatchNormalization() 71 | 72 | def call(self, inputs, training=None, **kwargs): 73 | x = self.a(inputs) 74 | x = self.a_bn(x, training=training) 75 | x = tf.keras.activations.relu(x) 76 | x = self.b(x) 77 | x = self.b_bn(x, training=training) 78 | return x 79 | 80 | 81 | class ResBasicBlock(tf.keras.layers.Layer): 82 | def __init__(self, w_in, w_out, stride): 83 | super(ResBasicBlock, self).__init__() 84 | self.proj, self.bn = None, None 85 | if (w_in != w_out) or (stride != 1): 86 | self.proj = tf.keras.layers.Conv2D(filters=w_out, kernel_size=1, strides=stride, padding="same") 87 | self.bn = tf.keras.layers.BatchNormalization() 88 | self.f = BasicTransform(w_out, stride) 89 | 90 | def call(self, inputs, training=None, **kwargs): 91 | if self.proj: 92 | x = self.proj(inputs) 93 | x = self.bn(x, training=training) 94 | else: 95 | x = inputs 96 | return tf.nn.relu(x + self.f(inputs, training=training)) 97 | 98 | 99 | class BottleneckTransform(tf.keras.layers.Layer): 100 | def __init__(self, w_in, w_out, stride, params): 101 | super(BottleneckTransform, self).__init__() 102 | w_b = int(round(w_out * params["bot_mul"])) 103 | w_se = int(round(w_in * params["se_r"])) 104 | groups = w_b // params["group_w"] 105 | 106 | self.a = tf.keras.layers.Conv2D(filters=w_b, kernel_size=1, strides=1, padding="same") 107 | self.a_bn = tf.keras.layers.BatchNormalization() 108 | # self.b = tf.keras.layers.Conv2D(filters=w_b, kernel_size=3, strides=stride, padding="same", groups=groups) 109 | self.b = get_group_conv(in_channels=w_b, out_channels=w_b, kernel_size=3, strides=stride, padding="same", groups=groups) 110 | self.b_bn = tf.keras.layers.BatchNormalization() 111 | self.se = SE(w_b, w_se) if w_se else None 112 | self.c = tf.keras.layers.Conv2D(filters=w_out, kernel_size=1, strides=1, padding="same") 113 | self.c_bn = tf.keras.layers.BatchNormalization() 114 | 115 | def call(self, inputs, training=None, **kwargs): 116 | x = self.a(inputs) 117 | x = self.a_bn(x, training=training) 118 | x = tf.nn.relu(x) 119 | x = self.b(x) 120 | x = self.b_bn(x, training=training) 121 | x = tf.nn.relu(x) 122 | x = self.se(x) 123 | x = self.c(x) 124 | x = self.c_bn(x, training=training) 125 | return x 126 | 127 | 128 | class ResBottleneckBlock(tf.keras.layers.Layer): 129 | def __init__(self, w_in, w_out, stride, params): 130 | super(ResBottleneckBlock, self).__init__() 131 | self.proj, self.bn = None, None 132 | if (w_in != w_out) or (stride != 1): 133 | self.proj = tf.keras.layers.Conv2D(filters=w_out, kernel_size=1, strides=stride, padding="same") 134 | self.bn = tf.keras.layers.BatchNormalization() 135 | self.f = BottleneckTransform(w_in, w_out, stride, params) 136 | 137 | def call(self, inputs, training=None, **kwargs): 138 | if self.proj: 139 | x = self.proj(inputs) 140 | x = self.bn(x, training=training) 141 | else: 142 | x = inputs 143 | return tf.nn.relu(x + self.f(inputs, training=training)) 144 | 145 | 146 | class ResStemCifar(tf.keras.layers.Layer): 147 | def __init__(self, w_out): 148 | super(ResStemCifar, self).__init__() 149 | self.conv = tf.keras.layers.Conv2D(filters=w_out, kernel_size=3, strides=1, padding="same") 150 | self.bn = tf.keras.layers.BatchNormalization() 151 | 152 | def call(self, inputs, training=None, **kwargs): 153 | x = self.conv(inputs) 154 | x = self.bn(x, training=training) 155 | x = tf.nn.relu(x) 156 | return x 157 | 158 | 159 | class ResStem(tf.keras.layers.Layer): 160 | def __init__(self, w_out): 161 | super(ResStem, self).__init__() 162 | self.conv = tf.keras.layers.Conv2D(filters=w_out, kernel_size=7, strides=2, padding="same") 163 | self.bn = tf.keras.layers.BatchNormalization() 164 | 165 | self.pool = tf.keras.layers.MaxPool2D(pool_size=3, strides=2, padding="same") 166 | 167 | def call(self, inputs, training=None, **kwargs): 168 | x = self.conv(inputs) 169 | x = self.bn(x, training=training) 170 | x = tf.nn.relu(x) 171 | x = self.pool(x) 172 | return x 173 | 174 | 175 | class SimpleStem(tf.keras.layers.Layer): 176 | def __init__(self, w_out): 177 | super(SimpleStem, self).__init__() 178 | self.conv = tf.keras.layers.Conv2D(filters=w_out, kernel_size=3, strides=2, padding="same") 179 | self.bn = tf.keras.layers.BatchNormalization() 180 | 181 | def call(self, inputs, training=None, **kwargs): 182 | x = self.conv(inputs) 183 | x = self.bn(x, training=training) 184 | x = tf.nn.relu(x) 185 | return x 186 | 187 | 188 | class AnyStage(tf.keras.layers.Layer): 189 | def __init__(self, w_in, w_out, stride, d, block_fun, params): 190 | super(AnyStage, self).__init__() 191 | self.layers = list() 192 | for _ in range(d): 193 | self.block = block_fun(w_in, w_out, stride, params) 194 | self.layers.append(self.block) 195 | stride, w_in = 1, w_out 196 | 197 | def call(self, inputs, training=None, **kwargs): 198 | x = inputs 199 | for l in self.layers: 200 | x = l(x, training=training) 201 | return x 202 | 203 | 204 | class AnyNet(tf.keras.Model): 205 | def __init__(self, params=None): 206 | super(AnyNet, self).__init__() 207 | p = AnyNet.get_params() if not params else params 208 | stem_fun = get_stem_fun(p["stem_type"]) 209 | block_fun = get_block_fun(p["block_type"]) 210 | self.stem = stem_fun(p["stem_w"]) 211 | prev_w = p["stem_w"] 212 | keys = ["depths", "widths", "strides", "bot_muls", "group_ws"] 213 | 214 | self.stages = list() 215 | 216 | for i, (d, w, s, b, g) in enumerate(zip(*[p[k] for k in keys])): 217 | params = {"bot_mul": b, "group_w": g, "se_r": p["se_r"]} 218 | stage = AnyStage(prev_w, w, s, d, block_fun, params) 219 | self.stages.append(stage) 220 | prev_w = w 221 | self.head = AnyHead(p["num_classes"]) 222 | 223 | @staticmethod 224 | def get_params(): 225 | nones = [None for _ in AnyNetCfg.depths] 226 | return { 227 | "stem_type": AnyNetCfg.stem_type, 228 | "stem_w": AnyNetCfg.stem_w, 229 | "block_type": AnyNetCfg.block_type, 230 | "depths": AnyNetCfg.depths, 231 | "widths": AnyNetCfg.widths, 232 | "strides": AnyNetCfg.strides, 233 | "bot_muls": AnyNetCfg.bot_muls if AnyNetCfg.bot_muls else nones, 234 | "group_ws": AnyNetCfg.group_ws if AnyNetCfg.group_ws else nones, 235 | "se_r": AnyNetCfg.se_r if AnyNetCfg.se_on else 0, 236 | "num_classes": NUM_CLASSES, 237 | } 238 | 239 | def call(self, inputs, training=None, mask=None): 240 | x = self.stem(inputs, training=training) 241 | for s in self.stages: 242 | x = s(x, training=training) 243 | x = self.head(x) 244 | return x -------------------------------------------------------------------------------- /models/RegNet/blocks.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import tensorflow_addons as tfa 3 | import numpy as np 4 | 5 | 6 | class SE(tf.keras.layers.Layer): 7 | def __init__(self, w_in, w_se): 8 | super(SE, self).__init__() 9 | self.avg_pool = tfa.layers.AdaptiveAveragePooling2D(output_size=(1, 1)) 10 | 11 | self.c1 = tf.keras.layers.Conv2D(filters=w_se, kernel_size=1, padding="same") 12 | self.c2 = tf.keras.layers.Conv2D(filters=w_in, kernel_size=1, padding="same") 13 | 14 | def call(self, inputs, **kwargs): 15 | x = self.avg_pool(inputs) 16 | x = tf.nn.relu(self.c1(x)) 17 | x = tf.nn.sigmoid(self.c2(x)) 18 | return x 19 | 20 | 21 | def adjust_block_compatibility(ws, bs, gs): 22 | """Adjusts the compatibility of widths, bottlenecks, and groups.""" 23 | assert len(ws) == len(bs) == len(gs) 24 | assert all(w > 0 and b > 0 and g > 0 for w, b, g in zip(ws, bs, gs)) 25 | vs = [int(max(1, w * b)) for w, b in zip(ws, bs)] 26 | gs = [int(min(g, v)) for g, v in zip(gs, vs)] 27 | ms = [np.lcm(g, b) if b > 1 else g for g, b in zip(gs, bs)] 28 | vs = [max(m, int(round(v / m) * m)) for v, m in zip(vs, ms)] 29 | ws = [int(v / b) for v, b in zip(vs, bs)] 30 | assert all(w * b % g == 0 for w, b, g in zip(ws, bs, gs)) 31 | return ws, bs, gs -------------------------------------------------------------------------------- /models/RegNet/regnet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import models.RegNet.blocks as bk 3 | 4 | from models.RegNet.anynet import AnyNet 5 | from anynet_cfg import RegNetCfg 6 | from configuration import NUM_CLASSES 7 | 8 | 9 | def generate_regnet(w_a, w_0, w_m, d, q=8): 10 | """Generates per stage widths and depths from RegNet parameters.""" 11 | assert w_a >= 0 and w_0 > 0 and w_m > 1 and w_0 % q == 0 12 | # Generate continuous per-block ws 13 | ws_cont = np.arange(d) * w_a + w_0 14 | # Generate quantized per-block ws 15 | ks = np.round(np.log(ws_cont / w_0) / np.log(w_m)) 16 | ws_all = w_0 * np.power(w_m, ks) 17 | ws_all = np.round(np.divide(ws_all, q)).astype(int) * q 18 | # Generate per stage ws and ds (assumes ws_all are sorted) 19 | ws, ds = np.unique(ws_all, return_counts=True) 20 | # Compute number of actual stages and total possible stages 21 | num_stages, total_stages = len(ws), ks.max() + 1 22 | # Convert numpy arrays to lists and return 23 | ws, ds, ws_all, ws_cont = (x.tolist() for x in (ws, ds, ws_all, ws_cont)) 24 | return ws, ds, num_stages, total_stages, ws_all, ws_cont 25 | 26 | 27 | class RegNet(AnyNet): 28 | @staticmethod 29 | def get_params(): 30 | w_a, w_0, w_m, d = RegNetCfg.wa, RegNetCfg.w0, RegNetCfg.wm, RegNetCfg.depth 31 | ws, ds = generate_regnet(w_a, w_0, w_m, d)[0:2] 32 | ss = [RegNetCfg.stride for _ in ws] 33 | bs = [RegNetCfg.bot_mul for _ in ws] 34 | gs = [RegNetCfg.group_w for _ in ws] 35 | ws, bs, gs = bk.adjust_block_compatibility(ws, bs, gs) 36 | 37 | return { 38 | "stem_type": RegNetCfg.stem_type, 39 | "stem_w": RegNetCfg.stem_w, 40 | "block_type": RegNetCfg.block_type, 41 | "depths": ds, 42 | "widths": ws, 43 | "strides": ss, 44 | "bot_muls": bs, 45 | "group_ws": gs, 46 | "se_r": RegNetCfg.se_r if RegNetCfg.se_on else 0, 47 | "num_classes": NUM_CLASSES, 48 | } 49 | 50 | def __init__(self): 51 | params = RegNet.get_params() 52 | super(RegNet, self).__init__(params) 53 | 54 | def __repr__(self): 55 | return "RegNet" 56 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | from . import mobilenet_v1, mobilenet_v2, mobilenet_v3_large, mobilenet_v3_small, \ 2 | efficientnet, resnext, inception_v4, inception_resnet_v1, inception_resnet_v2, \ 3 | se_resnet, squeezenet, densenet, shufflenet_v2, resnet, se_resnext 4 | from .RegNet import regnet 5 | 6 | model_list = [ 7 | mobilenet_v1.MobileNetV1(), mobilenet_v2.MobileNetV2(), mobilenet_v3_large.MobileNetV3Large(), 8 | mobilenet_v3_small.MobileNetV3Small(), 9 | efficientnet.efficient_net_b0(), efficientnet.efficient_net_b1(), 10 | efficientnet.efficient_net_b2(), efficientnet.efficient_net_b3(), efficientnet.efficient_net_b4(), 11 | efficientnet.efficient_net_b5(), efficientnet.efficient_net_b6(), efficientnet.efficient_net_b7(), 12 | resnext.ResNeXt50(), resnext.ResNeXt101(), 13 | inception_v4.InceptionV4(), 14 | inception_resnet_v1.InceptionResNetV1(), 15 | inception_resnet_v2.InceptionResNetV2(), 16 | se_resnet.se_resnet_50(), se_resnet.se_resnet_101(), se_resnet.se_resnet_152(), 17 | se_resnext.SEResNeXt50(), se_resnext.SEResNeXt101(), 18 | resnet.resnet_18(), resnet.resnet_34(), resnet.resnet_50(), resnet.resnet_101(), resnet.resnet_152(), 19 | shufflenet_v2.shufflenet_0_5x(), shufflenet_v2.shufflenet_1_0x(), shufflenet_v2.shufflenet_1_5x(), 20 | shufflenet_v2.shufflenet_2_0x(), 21 | regnet.RegNet() 22 | ] 23 | 24 | 25 | def get_model2idx_dict(): 26 | return dict((v, k) for k, v in enumerate(model_list)) 27 | 28 | 29 | def get_model(idx): 30 | return model_list[idx] 31 | -------------------------------------------------------------------------------- /models/densenet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from configuration import NUM_CLASSES 3 | 4 | 5 | class BottleNeck(tf.keras.layers.Layer): 6 | def __init__(self, growth_rate, drop_rate): 7 | super(BottleNeck, self).__init__() 8 | self.bn1 = tf.keras.layers.BatchNormalization() 9 | self.conv1 = tf.keras.layers.Conv2D(filters=4 * growth_rate, 10 | kernel_size=(1, 1), 11 | strides=1, 12 | padding="same") 13 | self.bn2 = tf.keras.layers.BatchNormalization() 14 | self.conv2 = tf.keras.layers.Conv2D(filters=growth_rate, 15 | kernel_size=(3, 3), 16 | strides=1, 17 | padding="same") 18 | self.dropout = tf.keras.layers.Dropout(rate=drop_rate) 19 | 20 | def call(self, inputs, training=None, **kwargs): 21 | x = self.bn1(inputs, training=training) 22 | x = tf.nn.relu(x) 23 | x = self.conv1(x) 24 | x = self.bn2(x, training=training) 25 | x = tf.nn.relu(x) 26 | x = self.conv2(x) 27 | x = self.dropout(x, training=training) 28 | return x 29 | 30 | 31 | class DenseBlock(tf.keras.layers.Layer): 32 | def __init__(self, num_layers, growth_rate, drop_rate): 33 | super(DenseBlock, self).__init__() 34 | self.num_layers = num_layers 35 | self.growth_rate = growth_rate 36 | self.drop_rate = drop_rate 37 | self.features_list = [] 38 | self.bottle_necks = [] 39 | for i in range(self.num_layers): 40 | self.bottle_necks.append(BottleNeck(growth_rate=self.growth_rate, drop_rate=self.drop_rate)) 41 | 42 | def call(self, inputs, training=None, **kwargs): 43 | self.features_list.append(inputs) 44 | x = inputs 45 | for i in range(self.num_layers): 46 | y = self.bottle_necks[i](x, training=training) 47 | self.features_list.append(y) 48 | x = tf.concat(self.features_list, axis=-1) 49 | self.features_list.clear() 50 | return x 51 | 52 | 53 | class TransitionLayer(tf.keras.layers.Layer): 54 | def __init__(self, out_channels): 55 | super(TransitionLayer, self).__init__() 56 | self.bn = tf.keras.layers.BatchNormalization() 57 | self.conv = tf.keras.layers.Conv2D(filters=out_channels, 58 | kernel_size=(1, 1), 59 | strides=1, 60 | padding="same") 61 | self.pool = tf.keras.layers.MaxPool2D(pool_size=(2, 2), 62 | strides=2, 63 | padding="same") 64 | 65 | def call(self, inputs, training=None, **kwargs): 66 | x = self.bn(inputs, training=training) 67 | x = tf.nn.relu(x) 68 | x = self.conv(x) 69 | x = self.pool(x) 70 | return x 71 | 72 | 73 | class DenseNet(tf.keras.Model): 74 | def __init__(self, num_init_features, growth_rate, block_layers, compression_rate, drop_rate): 75 | super(DenseNet, self).__init__() 76 | self.conv = tf.keras.layers.Conv2D(filters=num_init_features, 77 | kernel_size=(7, 7), 78 | strides=2, 79 | padding="same") 80 | self.bn = tf.keras.layers.BatchNormalization() 81 | self.pool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 82 | strides=2, 83 | padding="same") 84 | self.num_channels = num_init_features 85 | self.dense_block_1 = DenseBlock(num_layers=block_layers[0], growth_rate=growth_rate, drop_rate=drop_rate) 86 | self.num_channels += growth_rate * block_layers[0] 87 | self.num_channels = compression_rate * self.num_channels 88 | self.transition_1 = TransitionLayer(out_channels=int(self.num_channels)) 89 | self.dense_block_2 = DenseBlock(num_layers=block_layers[1], growth_rate=growth_rate, drop_rate=drop_rate) 90 | self.num_channels += growth_rate * block_layers[1] 91 | self.num_channels = compression_rate * self.num_channels 92 | self.transition_2 = TransitionLayer(out_channels=int(self.num_channels)) 93 | self.dense_block_3 = DenseBlock(num_layers=block_layers[2], growth_rate=growth_rate, drop_rate=drop_rate) 94 | self.num_channels += growth_rate * block_layers[2] 95 | self.num_channels = compression_rate * self.num_channels 96 | self.transition_3 = TransitionLayer(out_channels=int(self.num_channels)) 97 | self.dense_block_4 = DenseBlock(num_layers=block_layers[3], growth_rate=growth_rate, drop_rate=drop_rate) 98 | 99 | self.avgpool = tf.keras.layers.GlobalAveragePooling2D() 100 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, 101 | activation=tf.keras.activations.softmax) 102 | 103 | def call(self, inputs, training=None, mask=None): 104 | x = self.conv(inputs) 105 | x = self.bn(x, training=training) 106 | x = tf.nn.relu(x) 107 | x = self.pool(x) 108 | 109 | x = self.dense_block_1(x, training=training) 110 | x = self.transition_1(x, training=training) 111 | x = self.dense_block_2(x, training=training) 112 | x = self.transition_2(x, training=training) 113 | x = self.dense_block_3(x, training=training) 114 | x = self.transition_3(x, training=training) 115 | x = self.dense_block_4(x, training=training) 116 | 117 | x = self.avgpool(x) 118 | x = self.fc(x) 119 | 120 | return x 121 | 122 | 123 | def densenet_121(): 124 | return DenseNet(num_init_features=64, growth_rate=32, block_layers=[6, 12, 24, 16], compression_rate=0.5, drop_rate=0.5) 125 | 126 | 127 | def densenet_169(): 128 | return DenseNet(num_init_features=64, growth_rate=32, block_layers=[6, 12, 32, 32], compression_rate=0.5, drop_rate=0.5) 129 | 130 | 131 | def densenet_201(): 132 | return DenseNet(num_init_features=64, growth_rate=32, block_layers=[6, 12, 48, 32], compression_rate=0.5, drop_rate=0.5) 133 | 134 | 135 | def densenet_264(): 136 | return DenseNet(num_init_features=64, growth_rate=32, block_layers=[6, 12, 64, 48], compression_rate=0.5, drop_rate=0.5) 137 | -------------------------------------------------------------------------------- /models/efficientnet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import math 3 | from configuration import NUM_CLASSES 4 | 5 | 6 | def round_filters(filters, multiplier): 7 | depth_divisor = 8 8 | min_depth = None 9 | min_depth = min_depth or depth_divisor 10 | filters = filters * multiplier 11 | new_filters = max(min_depth, int(filters + depth_divisor / 2) // depth_divisor * depth_divisor) 12 | if new_filters < 0.9 * filters: 13 | new_filters += depth_divisor 14 | return int(new_filters) 15 | 16 | 17 | def round_repeats(repeats, multiplier): 18 | if not multiplier: 19 | return repeats 20 | return int(math.ceil(multiplier * repeats)) 21 | 22 | 23 | class SEBlock(tf.keras.layers.Layer): 24 | def __init__(self, input_channels, ratio=0.25): 25 | super(SEBlock, self).__init__() 26 | self.num_reduced_filters = max(1, int(input_channels * ratio)) 27 | self.pool = tf.keras.layers.GlobalAveragePooling2D() 28 | self.reduce_conv = tf.keras.layers.Conv2D(filters=self.num_reduced_filters, 29 | kernel_size=(1, 1), 30 | strides=1, 31 | padding="same") 32 | self.expand_conv = tf.keras.layers.Conv2D(filters=input_channels, 33 | kernel_size=(1, 1), 34 | strides=1, 35 | padding="same") 36 | 37 | def call(self, inputs, **kwargs): 38 | branch = self.pool(inputs) 39 | branch = tf.expand_dims(input=branch, axis=1) 40 | branch = tf.expand_dims(input=branch, axis=1) 41 | branch = self.reduce_conv(branch) 42 | branch = tf.nn.swish(branch) 43 | branch = self.expand_conv(branch) 44 | branch = tf.nn.sigmoid(branch) 45 | output = inputs * branch 46 | return output 47 | 48 | 49 | class MBConv(tf.keras.layers.Layer): 50 | def __init__(self, in_channels, out_channels, expansion_factor, stride, k, drop_connect_rate): 51 | super(MBConv, self).__init__() 52 | self.in_channels = in_channels 53 | self.out_channels = out_channels 54 | self.stride = stride 55 | self.drop_connect_rate = drop_connect_rate 56 | self.conv1 = tf.keras.layers.Conv2D(filters=in_channels * expansion_factor, 57 | kernel_size=(1, 1), 58 | strides=1, 59 | padding="same", 60 | use_bias=False) 61 | self.bn1 = tf.keras.layers.BatchNormalization() 62 | self.dwconv = tf.keras.layers.DepthwiseConv2D(kernel_size=(k, k), 63 | strides=stride, 64 | padding="same", 65 | use_bias=False) 66 | self.bn2 = tf.keras.layers.BatchNormalization() 67 | self.se = SEBlock(input_channels=in_channels * expansion_factor) 68 | self.conv2 = tf.keras.layers.Conv2D(filters=out_channels, 69 | kernel_size=(1, 1), 70 | strides=1, 71 | padding="same", 72 | use_bias=False) 73 | self.bn3 = tf.keras.layers.BatchNormalization() 74 | self.dropout = tf.keras.layers.Dropout(rate=drop_connect_rate) 75 | 76 | def call(self, inputs, training=None, **kwargs): 77 | x = self.conv1(inputs) 78 | x = self.bn1(x, training=training) 79 | x = tf.nn.swish(x) 80 | x = self.dwconv(x) 81 | x = self.bn2(x, training=training) 82 | x = self.se(x) 83 | x = tf.nn.swish(x) 84 | x = self.conv2(x) 85 | x = self.bn3(x, training=training) 86 | if self.stride == 1 and self.in_channels == self.out_channels: 87 | if self.drop_connect_rate: 88 | x = self.dropout(x, training=training) 89 | x = tf.keras.layers.add([x, inputs]) 90 | return x 91 | 92 | 93 | def build_mbconv_block(in_channels, out_channels, layers, stride, expansion_factor, k, drop_connect_rate): 94 | block = tf.keras.Sequential() 95 | for i in range(layers): 96 | if i == 0: 97 | block.add(MBConv(in_channels=in_channels, 98 | out_channels=out_channels, 99 | expansion_factor=expansion_factor, 100 | stride=stride, 101 | k=k, 102 | drop_connect_rate=drop_connect_rate)) 103 | else: 104 | block.add(MBConv(in_channels=out_channels, 105 | out_channels=out_channels, 106 | expansion_factor=expansion_factor, 107 | stride=1, 108 | k=k, 109 | drop_connect_rate=drop_connect_rate)) 110 | return block 111 | 112 | 113 | class EfficientNet(tf.keras.Model): 114 | def __init__(self, width_coefficient, depth_coefficient, dropout_rate, drop_connect_rate=0.2, name="B0"): 115 | super(EfficientNet, self).__init__() 116 | self.model_name = name 117 | 118 | self.conv1 = tf.keras.layers.Conv2D(filters=round_filters(32, width_coefficient), 119 | kernel_size=(3, 3), 120 | strides=2, 121 | padding="same", 122 | use_bias=False) 123 | self.bn1 = tf.keras.layers.BatchNormalization() 124 | self.block1 = build_mbconv_block(in_channels=round_filters(32, width_coefficient), 125 | out_channels=round_filters(16, width_coefficient), 126 | layers=round_repeats(1, depth_coefficient), 127 | stride=1, 128 | expansion_factor=1, k=3, drop_connect_rate=drop_connect_rate) 129 | self.block2 = build_mbconv_block(in_channels=round_filters(16, width_coefficient), 130 | out_channels=round_filters(24, width_coefficient), 131 | layers=round_repeats(2, depth_coefficient), 132 | stride=2, 133 | expansion_factor=6, k=3, drop_connect_rate=drop_connect_rate) 134 | self.block3 = build_mbconv_block(in_channels=round_filters(24, width_coefficient), 135 | out_channels=round_filters(40, width_coefficient), 136 | layers=round_repeats(2, depth_coefficient), 137 | stride=2, 138 | expansion_factor=6, k=5, drop_connect_rate=drop_connect_rate) 139 | self.block4 = build_mbconv_block(in_channels=round_filters(40, width_coefficient), 140 | out_channels=round_filters(80, width_coefficient), 141 | layers=round_repeats(3, depth_coefficient), 142 | stride=2, 143 | expansion_factor=6, k=3, drop_connect_rate=drop_connect_rate) 144 | self.block5 = build_mbconv_block(in_channels=round_filters(80, width_coefficient), 145 | out_channels=round_filters(112, width_coefficient), 146 | layers=round_repeats(3, depth_coefficient), 147 | stride=1, 148 | expansion_factor=6, k=5, drop_connect_rate=drop_connect_rate) 149 | self.block6 = build_mbconv_block(in_channels=round_filters(112, width_coefficient), 150 | out_channels=round_filters(192, width_coefficient), 151 | layers=round_repeats(4, depth_coefficient), 152 | stride=2, 153 | expansion_factor=6, k=5, drop_connect_rate=drop_connect_rate) 154 | self.block7 = build_mbconv_block(in_channels=round_filters(192, width_coefficient), 155 | out_channels=round_filters(320, width_coefficient), 156 | layers=round_repeats(1, depth_coefficient), 157 | stride=1, 158 | expansion_factor=6, k=3, drop_connect_rate=drop_connect_rate) 159 | 160 | self.conv2 = tf.keras.layers.Conv2D(filters=round_filters(1280, width_coefficient), 161 | kernel_size=(1, 1), 162 | strides=1, 163 | padding="same", 164 | use_bias=False) 165 | self.bn2 = tf.keras.layers.BatchNormalization() 166 | self.pool = tf.keras.layers.GlobalAveragePooling2D() 167 | self.dropout = tf.keras.layers.Dropout(rate=dropout_rate) 168 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, 169 | activation=tf.keras.activations.softmax) 170 | 171 | def call(self, inputs, training=None, mask=None): 172 | x = self.conv1(inputs) 173 | x = self.bn1(x, training=training) 174 | x = tf.nn.swish(x) 175 | 176 | x = self.block1(x) 177 | x = self.block2(x) 178 | x = self.block3(x) 179 | x = self.block4(x) 180 | x = self.block5(x) 181 | x = self.block6(x) 182 | x = self.block7(x) 183 | 184 | x = self.conv2(x) 185 | x = self.bn2(x, training=training) 186 | x = tf.nn.swish(x) 187 | x = self.pool(x) 188 | x = self.dropout(x, training=training) 189 | x = self.fc(x) 190 | 191 | return x 192 | 193 | def __repr__(self): 194 | return "EfficientNet-{}".format(self.model_name) 195 | 196 | 197 | def get_efficient_net(width_coefficient, depth_coefficient, resolution, dropout_rate, name): 198 | net = EfficientNet(width_coefficient=width_coefficient, 199 | depth_coefficient=depth_coefficient, 200 | dropout_rate=dropout_rate, 201 | name=name) 202 | 203 | return net 204 | 205 | 206 | def efficient_net_b0(): 207 | return get_efficient_net(1.0, 1.0, 224, 0.2, "B0") 208 | 209 | 210 | def efficient_net_b1(): 211 | return get_efficient_net(1.0, 1.1, 240, 0.2, "B1") 212 | 213 | 214 | def efficient_net_b2(): 215 | return get_efficient_net(1.1, 1.2, 260, 0.3, "B2") 216 | 217 | 218 | def efficient_net_b3(): 219 | return get_efficient_net(1.2, 1.4, 300, 0.3, "B3") 220 | 221 | 222 | def efficient_net_b4(): 223 | return get_efficient_net(1.4, 1.8, 380, 0.4, "B4") 224 | 225 | 226 | def efficient_net_b5(): 227 | return get_efficient_net(1.6, 2.2, 456, 0.4, "B5") 228 | 229 | 230 | def efficient_net_b6(): 231 | return get_efficient_net(1.8, 2.6, 528, 0.5, "B6") 232 | 233 | 234 | def efficient_net_b7(): 235 | return get_efficient_net(2.0, 3.1, 600, 0.5, "B7") 236 | -------------------------------------------------------------------------------- /models/group_convolution.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | 4 | from tensorflow.keras import initializers, regularizers, constraints 5 | from tensorflow.keras import activations 6 | from configuration import DEVICE 7 | 8 | 9 | def get_group_conv(in_channels, 10 | out_channels, 11 | kernel_size, 12 | strides=(1, 1), 13 | padding='valid', 14 | groups=1): 15 | if DEVICE == "cpu": 16 | return GroupConv2D(input_channels=in_channels, 17 | output_channels=out_channels, 18 | kernel_size=kernel_size, 19 | strides=strides, 20 | padding=padding, 21 | groups=groups) 22 | elif DEVICE == "gpu": 23 | return tf.keras.layers.Conv2D(filters=out_channels, 24 | kernel_size=kernel_size, 25 | strides=strides, 26 | padding=padding, 27 | groups=groups) 28 | else: 29 | raise ValueError("Attribute 'DEVICE' must be 'cpu' or 'gpu'.") 30 | 31 | 32 | class GroupConv2D(tf.keras.layers.Layer): 33 | def __init__(self, 34 | input_channels, 35 | output_channels, 36 | kernel_size, 37 | strides=(1, 1), 38 | padding='valid', 39 | data_format=None, 40 | dilation_rate=(1, 1), 41 | activation=None, 42 | groups=1, 43 | use_bias=True, 44 | kernel_initializer='glorot_uniform', 45 | bias_initializer='zeros', 46 | kernel_regularizer=None, 47 | bias_regularizer=None, 48 | activity_regularizer=None, 49 | kernel_constraint=None, 50 | bias_constraint=None, 51 | **kwargs): 52 | super(GroupConv2D, self).__init__() 53 | 54 | if not input_channels % groups == 0: 55 | raise ValueError("The value of input_channels must be divisible by the value of groups.") 56 | if not output_channels % groups == 0: 57 | raise ValueError("The value of output_channels must be divisible by the value of groups.") 58 | 59 | self.input_channels = input_channels 60 | self.output_channels = output_channels 61 | self.kernel_size = kernel_size 62 | self.strides = strides 63 | self.padding = padding 64 | self.data_format = data_format 65 | self.dilation_rate = dilation_rate 66 | self.activation = activation 67 | self.groups = groups 68 | self.use_bias = use_bias 69 | self.kernel_initializer = kernel_initializer 70 | self.bias_initializer = bias_initializer 71 | self.kernel_regularizer = kernel_regularizer 72 | self.bias_regularizer = bias_regularizer 73 | self.activity_regularizer = activity_regularizer 74 | self.kernel_constraint = kernel_constraint 75 | self.bias_constraint = bias_constraint 76 | 77 | self.group_in_num = input_channels // groups 78 | self.group_out_num = output_channels // groups 79 | self.conv_list = [] 80 | for i in range(self.groups): 81 | self.conv_list.append(tf.keras.layers.Conv2D(filters=self.group_out_num, 82 | kernel_size=kernel_size, 83 | strides=strides, 84 | padding=padding, 85 | data_format=data_format, 86 | dilation_rate=dilation_rate, 87 | activation=activations.get(activation), 88 | use_bias=use_bias, 89 | kernel_initializer=initializers.get(kernel_initializer), 90 | bias_initializer=initializers.get(bias_initializer), 91 | kernel_regularizer=regularizers.get(kernel_regularizer), 92 | bias_regularizer=regularizers.get(bias_regularizer), 93 | activity_regularizer=regularizers.get(activity_regularizer), 94 | kernel_constraint=constraints.get(kernel_constraint), 95 | bias_constraint=constraints.get(bias_constraint), 96 | **kwargs)) 97 | 98 | def call(self, inputs, **kwargs): 99 | feature_map_list = [] 100 | for i in range(self.groups): 101 | x_i = self.conv_list[i](inputs[:, :, :, i*self.group_in_num: (i + 1) * self.group_in_num]) 102 | feature_map_list.append(x_i) 103 | out = tf.concat(feature_map_list, axis=-1) 104 | return out 105 | 106 | def get_config(self): 107 | config = { 108 | "input_channels": self.input_channels, 109 | "output_channels": self.output_channels, 110 | "kernel_size": self.kernel_size, 111 | "strides": self.strides, 112 | "padding": self.padding, 113 | "data_format": self.data_format, 114 | "dilation_rate": self.dilation_rate, 115 | "activation": activations.serialize(self.activation), 116 | "groups": self.groups, 117 | "use_bias": self.use_bias, 118 | "kernel_initializer": initializers.serialize(self.kernel_initializer), 119 | "bias_initializer": initializers.serialize(self.bias_initializer), 120 | "kernel_regularizer": regularizers.serialize(self.kernel_regularizer), 121 | "bias_regularizer": regularizers.serialize(self.bias_regularizer), 122 | "activity_regularizer": regularizers.serialize(self.activity_regularizer), 123 | "kernel_constraint": constraints.serialize(self.kernel_constraint), 124 | "bias_constraint": constraints.serialize(self.bias_constraint) 125 | } 126 | base_config = super(GroupConv2D, self).get_config() 127 | return {**base_config, **config} 128 | 129 | 130 | class GroupConv2DTranspose(tf.keras.layers.Layer): 131 | def __init__(self, 132 | input_channels, 133 | output_channels, 134 | kernel_size, 135 | strides=(1, 1), 136 | padding='valid', 137 | output_padding=None, 138 | data_format=None, 139 | dilation_rate=(1, 1), 140 | activation=None, 141 | groups=1, 142 | use_bias=True, 143 | kernel_initializer='glorot_uniform', 144 | bias_initializer='zeros', 145 | kernel_regularizer=None, 146 | bias_regularizer=None, 147 | activity_regularizer=None, 148 | kernel_constraint=None, 149 | bias_constraint=None, 150 | **kwargs 151 | ): 152 | super(GroupConv2DTranspose, self).__init__() 153 | 154 | if not input_channels % groups == 0: 155 | raise ValueError("The value of input_channels must be divisible by the value of groups.") 156 | if not output_channels % groups == 0: 157 | raise ValueError("The value of output_channels must be divisible by the value of groups.") 158 | 159 | self.input_channels = input_channels 160 | self.output_channels = output_channels 161 | self.kernel_size = kernel_size 162 | self.strides = strides 163 | self.padding = padding 164 | self.output_padding = output_padding 165 | self.data_format = data_format 166 | self.dilation_rate = dilation_rate 167 | self.activation = activation 168 | self.groups = groups 169 | self.use_bias = use_bias 170 | self.kernel_initializer = kernel_initializer 171 | self.bias_initializer = bias_initializer 172 | self.kernel_regularizer = kernel_regularizer 173 | self.bias_regularizer = bias_regularizer 174 | self.activity_regularizer = activity_regularizer 175 | self.kernel_constraint = kernel_constraint 176 | self.bias_constraint = bias_constraint 177 | 178 | self.group_in_num = input_channels // groups 179 | self.group_out_num = output_channels // groups 180 | self.conv_list = [] 181 | for i in range(self.groups): 182 | self.conv_list.append(tf.keras.layers.Conv2DTranspose(filters=self.group_out_num, 183 | kernel_size=kernel_size, 184 | strides=strides, 185 | padding=padding, 186 | output_padding=output_padding, 187 | data_format=data_format, 188 | dilation_rate=dilation_rate, 189 | activation=activations.get(activation), 190 | use_bias=use_bias, 191 | kernel_initializer=initializers.get(kernel_initializer), 192 | bias_initializer=initializers.get(bias_initializer), 193 | kernel_regularizer=regularizers.get(kernel_regularizer), 194 | bias_regularizer=regularizers.get(bias_regularizer), 195 | activity_regularizer=regularizers.get(activity_regularizer), 196 | kernel_constraint=constraints.get(kernel_constraint), 197 | bias_constraint=constraints.get(bias_constraint), 198 | **kwargs)) 199 | 200 | def call(self, inputs, **kwargs): 201 | feature_map_list = [] 202 | for i in range(self.groups): 203 | x_i = self.conv_list[i](inputs[:, :, :, i*self.group_in_num: (i + 1) * self.group_in_num]) 204 | feature_map_list.append(x_i) 205 | out = tf.concat(feature_map_list, axis=-1) 206 | return out 207 | 208 | def get_config(self): 209 | config = { 210 | "input_channels": self.input_channels, 211 | "output_channels": self.output_channels, 212 | "kernel_size": self.kernel_size, 213 | "strides": self.strides, 214 | "padding": self.padding, 215 | "output_padding": self.output_padding, 216 | "data_format": self.data_format, 217 | "dilation_rate": self.dilation_rate, 218 | "activation": activations.serialize(self.activation), 219 | "groups": self.groups, 220 | "use_bias": self.use_bias, 221 | "kernel_initializer": initializers.serialize(self.kernel_initializer), 222 | "bias_initializer": initializers.serialize(self.bias_initializer), 223 | "kernel_regularizer": regularizers.serialize(self.kernel_regularizer), 224 | "bias_regularizer": regularizers.serialize(self.bias_regularizer), 225 | "activity_regularizer": regularizers.serialize(self.activity_regularizer), 226 | "kernel_constraint": constraints.serialize(self.kernel_constraint), 227 | "bias_constraint": constraints.serialize(self.bias_constraint) 228 | } 229 | base_config = super(GroupConv2DTranspose, self).get_config() 230 | return {**base_config, **config} 231 | -------------------------------------------------------------------------------- /models/inception_modules.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | 4 | class BasicConv2D(tf.keras.layers.Layer): 5 | def __init__(self, filters, kernel_size, strides, padding): 6 | super(BasicConv2D, self).__init__() 7 | self.conv = tf.keras.layers.Conv2D(filters=filters, 8 | kernel_size=kernel_size, 9 | strides=strides, 10 | padding=padding) 11 | self.bn = tf.keras.layers.BatchNormalization() 12 | 13 | def call(self, inputs, training=None, **kwargs): 14 | x = self.conv(inputs) 15 | x = self.bn(x, training=training) 16 | x = tf.nn.relu(x) 17 | 18 | return x 19 | 20 | 21 | class Conv2DLinear(tf.keras.layers.Layer): 22 | def __init__(self, filters, kernel_size, strides, padding): 23 | super(Conv2DLinear, self).__init__() 24 | self.conv = tf.keras.layers.Conv2D(filters=filters, 25 | kernel_size=kernel_size, 26 | strides=strides, 27 | padding=padding) 28 | self.bn = tf.keras.layers.BatchNormalization() 29 | 30 | def call(self, inputs, training=None, **kwargs): 31 | x = self.conv(inputs) 32 | x = self.bn(x, training=training) 33 | 34 | return x 35 | 36 | 37 | class Stem(tf.keras.layers.Layer): 38 | def __init__(self): 39 | super(Stem, self).__init__() 40 | self.conv1 = BasicConv2D(filters=32, 41 | kernel_size=(3, 3), 42 | strides=2, 43 | padding="valid") 44 | self.conv2 = BasicConv2D(filters=32, 45 | kernel_size=(3, 3), 46 | strides=1, 47 | padding="valid") 48 | self.conv3 = BasicConv2D(filters=64, 49 | kernel_size=(3, 3), 50 | strides=1, 51 | padding="same") 52 | self.b1_maxpool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 53 | strides=2, 54 | padding="valid") 55 | self.b2_conv = BasicConv2D(filters=96, 56 | kernel_size=(3, 3), 57 | strides=2, 58 | padding="valid") 59 | self.b3_conv1 = BasicConv2D(filters=64, 60 | kernel_size=(1, 1), 61 | strides=1, 62 | padding="same") 63 | self.b3_conv2 = BasicConv2D(filters=96, 64 | kernel_size=(3, 3), 65 | strides=1, 66 | padding="valid") 67 | self.b4_conv1 = BasicConv2D(filters=64, 68 | kernel_size=(1, 1), 69 | strides=1, 70 | padding="same") 71 | self.b4_conv2 = BasicConv2D(filters=64, 72 | kernel_size=(7, 1), 73 | strides=1, 74 | padding="same") 75 | self.b4_conv3 = BasicConv2D(filters=64, 76 | kernel_size=(1, 7), 77 | strides=1, 78 | padding="same") 79 | self.b4_conv4 = BasicConv2D(filters=96, 80 | kernel_size=(3, 3), 81 | strides=1, 82 | padding="valid") 83 | self.b5_conv = BasicConv2D(filters=192, 84 | kernel_size=(3, 3), 85 | strides=2, 86 | padding="valid") 87 | self.b6_maxpool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 88 | strides=2, 89 | padding="valid") 90 | 91 | def call(self, inputs, training=None, **kwargs): 92 | x = self.conv1(inputs, training=training) 93 | x = self.conv2(x, training=training) 94 | x = self.conv3(x, training=training) 95 | branch_1 = self.b1_maxpool(x) 96 | branch_2 = self.b2_conv(x, training=training) 97 | x = tf.concat(values=[branch_1, branch_2], axis=-1) 98 | branch_3 = self.b3_conv1(x, training=training) 99 | branch_3 = self.b3_conv2(branch_3, training=training) 100 | branch_4 = self.b4_conv1(x, training=training) 101 | branch_4 = self.b4_conv2(branch_4, training=training) 102 | branch_4 = self.b4_conv3(branch_4, training=training) 103 | branch_4 = self.b4_conv4(branch_4, training=training) 104 | x = tf.concat(values=[branch_3, branch_4], axis=-1) 105 | branch_5 = self.b5_conv(x, training=training) 106 | branch_6 = self.b6_maxpool(x, training=training) 107 | x = tf.concat(values=[branch_5, branch_6], axis=-1) 108 | 109 | return x 110 | 111 | 112 | class InceptionBlockA(tf.keras.layers.Layer): 113 | def __init__(self): 114 | super(InceptionBlockA, self).__init__() 115 | self.b1_pool = tf.keras.layers.AveragePooling2D(pool_size=(3, 3), 116 | strides=1, 117 | padding="same") 118 | self.b1_conv = BasicConv2D(filters=96, 119 | kernel_size=(1, 1), 120 | strides=1, 121 | padding="same") 122 | self.b2_conv = BasicConv2D(filters=96, 123 | kernel_size=(1, 1), 124 | strides=1, 125 | padding="same") 126 | self.b3_conv1 = BasicConv2D(filters=64, 127 | kernel_size=(1, 1), 128 | strides=1, 129 | padding="same") 130 | self.b3_conv2 = BasicConv2D(filters=96, 131 | kernel_size=(3, 3), 132 | strides=1, 133 | padding="same") 134 | self.b4_conv1 = BasicConv2D(filters=64, 135 | kernel_size=(1, 1), 136 | strides=1, 137 | padding="same") 138 | self.b4_conv2 = BasicConv2D(filters=96, 139 | kernel_size=(3, 3), 140 | strides=1, 141 | padding="same") 142 | self.b4_conv3 = BasicConv2D(filters=96, 143 | kernel_size=(3, 3), 144 | strides=1, 145 | padding="same") 146 | 147 | def call(self, inputs, training=None, **kwargs): 148 | b1 = self.b1_pool(inputs) 149 | b1 = self.b1_conv(b1, training=training) 150 | 151 | b2 = self.b2_conv(inputs, training=training) 152 | 153 | b3 = self.b3_conv1(inputs, training=training) 154 | b3 = self.b3_conv2(b3, training=training) 155 | 156 | b4 = self.b4_conv1(inputs, training=training) 157 | b4 = self.b4_conv2(b4, training=training) 158 | b4 = self.b4_conv3(b4, training=training) 159 | 160 | return tf.concat(values=[b1, b2, b3, b4], axis=-1) 161 | 162 | 163 | class ReductionA(tf.keras.layers.Layer): 164 | def __init__(self, k, l, m, n): 165 | super(ReductionA, self).__init__() 166 | self.b1_pool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 167 | strides=2, 168 | padding="valid") 169 | self.b2_conv = BasicConv2D(filters=n, 170 | kernel_size=(3, 3), 171 | strides=2, 172 | padding="valid") 173 | self.b3_conv1 = BasicConv2D(filters=k, 174 | kernel_size=(1, 1), 175 | strides=1, 176 | padding="same") 177 | self.b3_conv2 = BasicConv2D(filters=l, 178 | kernel_size=(3, 3), 179 | strides=1, 180 | padding="same") 181 | self.b3_conv3 = BasicConv2D(filters=m, 182 | kernel_size=(3, 3), 183 | strides=2, 184 | padding="valid") 185 | 186 | def call(self, inputs, training=None, **kwargs): 187 | b1 = self.b1_pool(inputs) 188 | 189 | b2 = self.b2_conv(inputs, training=training) 190 | 191 | b3 = self.b3_conv1(inputs, training=training) 192 | b3 = self.b3_conv2(b3, training=training) 193 | b3 = self.b3_conv3(b3, training=training) 194 | 195 | return tf.concat(values=[b1, b2, b3], axis=-1) 196 | 197 | 198 | class InceptionBlockB(tf.keras.layers.Layer): 199 | def __init__(self): 200 | super(InceptionBlockB, self).__init__() 201 | self.b1_pool = tf.keras.layers.AveragePooling2D(pool_size=(3, 3), 202 | strides=1, 203 | padding="same") 204 | self.b1_conv = BasicConv2D(filters=128, 205 | kernel_size=(1, 1), 206 | strides=1, 207 | padding="same") 208 | self.b2_conv = BasicConv2D(filters=384, 209 | kernel_size=(1, 1), 210 | strides=1, 211 | padding="same") 212 | self.b3_conv1 = BasicConv2D(filters=192, 213 | kernel_size=(1, 1), 214 | strides=1, 215 | padding="same") 216 | self.b3_conv2 = BasicConv2D(filters=224, 217 | kernel_size=(1, 7), 218 | strides=1, 219 | padding="same") 220 | self.b3_conv3 = BasicConv2D(filters=256, 221 | kernel_size=(1, 7), 222 | strides=1, 223 | padding="same") 224 | self.b4_conv1 = BasicConv2D(filters=192, 225 | kernel_size=(1, 1), 226 | strides=1, 227 | padding="same") 228 | self.b4_conv2 = BasicConv2D(filters=192, 229 | kernel_size=(1, 7), 230 | strides=1, 231 | padding="same") 232 | self.b4_conv3 = BasicConv2D(filters=224, 233 | kernel_size=(7, 1), 234 | strides=1, 235 | padding="same") 236 | self.b4_conv4 = BasicConv2D(filters=224, 237 | kernel_size=(1, 7), 238 | strides=1, 239 | padding="same") 240 | self.b4_conv5 = BasicConv2D(filters=256, 241 | kernel_size=(7, 1), 242 | strides=1, 243 | padding="same") 244 | 245 | def call(self, inputs, training=None, **kwargs): 246 | b1 = self.b1_pool(inputs) 247 | b1 = self.b1_conv(b1, training=training) 248 | 249 | b2 = self.b2_conv(inputs, training=training) 250 | 251 | b3 = self.b3_conv1(inputs, training=training) 252 | b3 = self.b3_conv2(b3, training=training) 253 | b3 = self.b3_conv3(b3, training=training) 254 | 255 | b4 = self.b4_conv1(inputs, training=training) 256 | b4 = self.b4_conv2(b4, training=training) 257 | b4 = self.b4_conv3(b4, training=training) 258 | b4 = self.b4_conv4(b4, training=training) 259 | b4 = self.b4_conv5(b4, training=training) 260 | 261 | return tf.concat(values=[b1, b2, b3, b4], axis=-1) 262 | 263 | 264 | class ReductionB(tf.keras.layers.Layer): 265 | def __init__(self): 266 | super(ReductionB, self).__init__() 267 | self.b1_pool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 268 | strides=2, 269 | padding="valid") 270 | self.b2_conv1 = BasicConv2D(filters=192, 271 | kernel_size=(1, 1), 272 | strides=1, 273 | padding="same") 274 | self.b2_conv2 = BasicConv2D(filters=192, 275 | kernel_size=(3, 3), 276 | strides=2, 277 | padding="valid") 278 | self.b3_conv1 = BasicConv2D(filters=256, 279 | kernel_size=(1, 1), 280 | strides=1, 281 | padding="same") 282 | self.b3_conv2 = BasicConv2D(filters=256, 283 | kernel_size=(1, 7), 284 | strides=1, 285 | padding="same") 286 | self.b3_conv3 = BasicConv2D(filters=320, 287 | kernel_size=(7, 1), 288 | strides=1, 289 | padding="same") 290 | self.b3_conv4 = BasicConv2D(filters=320, 291 | kernel_size=(3, 3), 292 | strides=2, 293 | padding="valid") 294 | 295 | def call(self, inputs, training=None, **kwargs): 296 | b1 = self.b1_pool(inputs) 297 | 298 | b2 = self.b2_conv1(inputs, training=training) 299 | b2 = self.b2_conv2(b2, training=training) 300 | 301 | b3 = self.b3_conv1(inputs, training=training) 302 | b3 = self.b3_conv2(b3, training=training) 303 | b3 = self.b3_conv3(b3, training=training) 304 | b3 = self.b3_conv4(b3, training=training) 305 | 306 | return tf.concat(values=[b1, b2, b3], axis=-1) 307 | 308 | 309 | class InceptionBlockC(tf.keras.layers.Layer): 310 | def __init__(self): 311 | super(InceptionBlockC, self).__init__() 312 | self.b1_pool = tf.keras.layers.AveragePooling2D(pool_size=(3, 3), 313 | strides=1, 314 | padding="same") 315 | self.b1_conv = BasicConv2D(filters=256, 316 | kernel_size=(1, 1), 317 | strides=1, 318 | padding="same") 319 | self.b2_conv = BasicConv2D(filters=256, 320 | kernel_size=(1, 1), 321 | strides=1, 322 | padding="same") 323 | self.b3_conv1 = BasicConv2D(filters=384, 324 | kernel_size=(1, 1), 325 | strides=1, 326 | padding="same") 327 | self.b3_conv2 = BasicConv2D(filters=256, 328 | kernel_size=(1, 3), 329 | strides=1, 330 | padding="same") 331 | self.b3_conv3 = BasicConv2D(filters=256, 332 | kernel_size=(3, 1), 333 | strides=1, 334 | padding="same") 335 | self.b4_conv1 = BasicConv2D(filters=384, 336 | kernel_size=(1, 1), 337 | strides=1, 338 | padding="same") 339 | self.b4_conv2 = BasicConv2D(filters=448, 340 | kernel_size=(1, 3), 341 | strides=1, 342 | padding="same") 343 | self.b4_conv3 = BasicConv2D(filters=512, 344 | kernel_size=(3, 1), 345 | strides=1, 346 | padding="same") 347 | self.b4_conv4 = BasicConv2D(filters=256, 348 | kernel_size=(3, 1), 349 | strides=1, 350 | padding="same") 351 | self.b4_conv5 = BasicConv2D(filters=256, 352 | kernel_size=(1, 3), 353 | strides=1, 354 | padding="same") 355 | 356 | def call(self, inputs, training=None, **kwargs): 357 | b1 = self.b1_pool(inputs) 358 | b1 = self.b1_conv(b1, training=training) 359 | 360 | b2 = self.b2_conv(inputs, training=training) 361 | 362 | b3 = self.b3_conv1(inputs, training=training) 363 | b3_1 = self.b3_conv2(b3, training=training) 364 | b3_2 = self.b3_conv3(b3, training=training) 365 | 366 | b4 = self.b4_conv1(inputs, training=training) 367 | b4 = self.b4_conv2(b4, training=training) 368 | b4 = self.b4_conv3(b4, training=training) 369 | b4_1 = self.b4_conv4(b4, training=training) 370 | b4_2 = self.b4_conv5(b4, training=training) 371 | 372 | return tf.concat(values=[b1, b2, b3_1, b3_2, b4_1, b4_2], axis=-1) 373 | -------------------------------------------------------------------------------- /models/inception_resnet_v1.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from models.inception_modules import BasicConv2D, Conv2DLinear, ReductionA 3 | from configuration import NUM_CLASSES 4 | 5 | 6 | class Stem(tf.keras.layers.Layer): 7 | def __init__(self): 8 | super(Stem, self).__init__() 9 | self.conv1 = BasicConv2D(filters=32, 10 | kernel_size=(3, 3), 11 | strides=2, 12 | padding="valid") 13 | self.conv2 = BasicConv2D(filters=32, 14 | kernel_size=(3, 3), 15 | strides=1, 16 | padding="valid") 17 | self.conv3 = BasicConv2D(filters=64, 18 | kernel_size=(3, 3), 19 | strides=1, 20 | padding="same") 21 | self.maxpool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 22 | strides=2, 23 | padding="valid") 24 | self.conv4 = BasicConv2D(filters=80, 25 | kernel_size=(1, 1), 26 | strides=1, 27 | padding="same") 28 | self.conv5 = BasicConv2D(filters=192, 29 | kernel_size=(3, 3), 30 | strides=1, 31 | padding="valid") 32 | self.conv6 = BasicConv2D(filters=256, 33 | kernel_size=(3, 3), 34 | strides=2, 35 | padding="valid") 36 | 37 | def call(self, inputs, training=None, **kwargs): 38 | x = self.conv1(inputs, training=training) 39 | x = self.conv2(x, training=training) 40 | x = self.conv3(x, training=training) 41 | x = self.maxpool(x) 42 | x = self.conv4(x, training=training) 43 | x = self.conv5(x, training=training) 44 | x = self.conv6(x, training=training) 45 | 46 | return x 47 | 48 | 49 | class InceptionResNetA(tf.keras.layers.Layer): 50 | def __init__(self): 51 | super(InceptionResNetA, self).__init__() 52 | self.b1_conv = BasicConv2D(filters=32, 53 | kernel_size=(1, 1), 54 | strides=1, 55 | padding="same") 56 | self.b2_conv1 = BasicConv2D(filters=32, 57 | kernel_size=(1, 1), 58 | strides=1, 59 | padding="same") 60 | self.b2_conv2 = BasicConv2D(filters=32, 61 | kernel_size=(3, 3), 62 | strides=1, 63 | padding="same") 64 | self.b3_conv1 = BasicConv2D(filters=32, 65 | kernel_size=(1, 1), 66 | strides=1, 67 | padding="same") 68 | self.b3_conv2 = BasicConv2D(filters=32, 69 | kernel_size=(3, 3), 70 | strides=1, 71 | padding="same") 72 | self.b3_conv3 = BasicConv2D(filters=32, 73 | kernel_size=(3, 3), 74 | strides=1, 75 | padding="same") 76 | self.conv = Conv2DLinear(filters=256, 77 | kernel_size=(1, 1), 78 | strides=1, 79 | padding="same") 80 | 81 | def call(self, inputs, training=None, **kwargs): 82 | b1 = self.b1_conv(inputs, training=training) 83 | b2 = self.b2_conv1(inputs, training=training) 84 | b2 = self.b2_conv2(b2, training=training) 85 | b3 = self.b3_conv1(inputs, training=training) 86 | b3 = self.b3_conv2(b3, training=training) 87 | b3 = self.b3_conv3(b3, training=training) 88 | 89 | x = tf.concat(values=[b1, b2, b3], axis=-1) 90 | x = self.conv(x, training=training) 91 | 92 | output = tf.keras.layers.add([x, inputs]) 93 | return tf.nn.relu(output) 94 | 95 | 96 | class InceptionResNetB(tf.keras.layers.Layer): 97 | def __init__(self): 98 | super(InceptionResNetB, self).__init__() 99 | self.b1_conv = BasicConv2D(filters=128, 100 | kernel_size=(1, 1), 101 | strides=1, 102 | padding="same") 103 | self.b2_conv1 = BasicConv2D(filters=128, 104 | kernel_size=(1, 1), 105 | strides=1, 106 | padding="same") 107 | self.b2_conv2 = BasicConv2D(filters=128, 108 | kernel_size=(1, 7), 109 | strides=1, 110 | padding="same") 111 | self.b2_conv3 = BasicConv2D(filters=128, 112 | kernel_size=(7, 1), 113 | strides=1, 114 | padding="same") 115 | self.conv = Conv2DLinear(filters=896, 116 | kernel_size=(1, 1), 117 | strides=1, 118 | padding="same") 119 | 120 | def call(self, inputs, training=None, **kwargs): 121 | b1 = self.b1_conv(inputs, training=training) 122 | 123 | b2 = self.b2_conv1(inputs, training=training) 124 | b2 = self.b2_conv2(b2, training=training) 125 | b2 = self.b2_conv3(b2, training=training) 126 | 127 | x = tf.concat(values=[b1, b2], axis=-1) 128 | x = self.conv(x, training=training) 129 | 130 | output = tf.keras.layers.add([x, inputs]) 131 | return tf.nn.relu(output) 132 | 133 | 134 | class ReductionB(tf.keras.layers.Layer): 135 | def __init__(self): 136 | super(ReductionB, self).__init__() 137 | self.b1_maxpool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 138 | strides=2, 139 | padding="valid") 140 | self.b2_conv1 = BasicConv2D(filters=256, 141 | kernel_size=(1, 1), 142 | strides=1, 143 | padding="same") 144 | self.b2_conv2 = BasicConv2D(filters=384, 145 | kernel_size=(3, 3), 146 | strides=2, 147 | padding="valid") 148 | self.b3_conv1 = BasicConv2D(filters=256, 149 | kernel_size=(1, 1), 150 | strides=1, 151 | padding="same") 152 | self.b3_conv2 = BasicConv2D(filters=256, 153 | kernel_size=(3, 3), 154 | strides=2, 155 | padding="valid") 156 | self.b4_conv1 = BasicConv2D(filters=256, 157 | kernel_size=(1, 1), 158 | strides=1, 159 | padding="same") 160 | self.b4_conv2 = BasicConv2D(filters=256, 161 | kernel_size=(3, 3), 162 | strides=1, 163 | padding="same") 164 | self.b4_conv3 = BasicConv2D(filters=256, 165 | kernel_size=(3, 3), 166 | strides=2, 167 | padding="valid") 168 | 169 | def call(self, inputs, training=None, **kwargs): 170 | b1 = self.b1_maxpool(inputs) 171 | 172 | b2 = self.b2_conv1(inputs, training=training) 173 | b2 = self.b2_conv2(b2, training=training) 174 | 175 | b3 = self.b3_conv1(inputs, training=training) 176 | b3 = self.b3_conv2(b3, training=training) 177 | 178 | b4 = self.b4_conv1(inputs, training=training) 179 | b4 = self.b4_conv2(b4, training=training) 180 | b4 = self.b4_conv3(b4, training=training) 181 | 182 | return tf.concat(values=[b1, b2, b3, b4], axis=-1) 183 | 184 | 185 | class InceptionResNetC(tf.keras.layers.Layer): 186 | def __init__(self): 187 | super(InceptionResNetC, self).__init__() 188 | self.b1_conv = BasicConv2D(filters=192, 189 | kernel_size=(1, 1), 190 | strides=1, 191 | padding="same") 192 | self.b2_conv1 = BasicConv2D(filters=192, 193 | kernel_size=(1, 1), 194 | strides=1, 195 | padding="same") 196 | self.b2_conv2 = BasicConv2D(filters=192, 197 | kernel_size=(1, 3), 198 | strides=1, 199 | padding="same") 200 | self.b2_conv3 = BasicConv2D(filters=192, 201 | kernel_size=(3, 1), 202 | strides=1, 203 | padding="same") 204 | self.conv = Conv2DLinear(filters=1792, 205 | kernel_size=(1, 1), 206 | strides=1, 207 | padding="same") 208 | 209 | def call(self, inputs, training=None, **kwargs): 210 | b1 = self.b1_conv(inputs, training=training) 211 | b2 = self.b2_conv1(inputs, training=training) 212 | b2 = self.b2_conv2(b2, training=training) 213 | b2 = self.b2_conv3(b2, training=training) 214 | 215 | x = tf.concat(values=[b1, b2], axis=-1) 216 | x = self.conv(x, training=training) 217 | 218 | output = tf.keras.layers.add([x, inputs]) 219 | return tf.nn.relu(output) 220 | 221 | 222 | def build_inception_resnet_a(n): 223 | block = tf.keras.Sequential() 224 | for _ in range(n): 225 | block.add(InceptionResNetA()) 226 | return block 227 | 228 | 229 | def build_inception_resnet_b(n): 230 | block = tf.keras.Sequential() 231 | for _ in range(n): 232 | block.add(InceptionResNetB()) 233 | return block 234 | 235 | 236 | def build_inception_resnet_c(n): 237 | block = tf.keras.Sequential() 238 | for _ in range(n): 239 | block.add(InceptionResNetC()) 240 | return block 241 | 242 | 243 | class InceptionResNetV1(tf.keras.Model): 244 | def __init__(self): 245 | super(InceptionResNetV1, self).__init__() 246 | self.stem = Stem() 247 | self.inception_resnet_a = build_inception_resnet_a(5) 248 | self.reduction_a = ReductionA(k=192, l=192, m=256, n=384) 249 | self.inception_resnet_b = build_inception_resnet_b(10) 250 | self.reduction_b = ReductionB() 251 | self.inception_resnet_c = build_inception_resnet_c(5) 252 | self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(8, 8)) 253 | self.dropout = tf.keras.layers.Dropout(rate=0.2) 254 | self.flat = tf.keras.layers.Flatten() 255 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, 256 | activation=tf.keras.activations.softmax) 257 | 258 | def call(self, inputs, training=None, mask=None): 259 | x = self.stem(inputs, training=training) 260 | x = self.inception_resnet_a(x, training=training) 261 | x = self.reduction_a(x, training=training) 262 | x = self.inception_resnet_b(x, training=training) 263 | x = self.reduction_b(x, training=training) 264 | x = self.inception_resnet_c(x, training=training) 265 | x = self.avgpool(x) 266 | x = self.dropout(x, training=training) 267 | x = self.flat(x) 268 | x = self.fc(x) 269 | 270 | return x 271 | 272 | def __repr__(self): 273 | return "InceptionResNetV1" -------------------------------------------------------------------------------- /models/inception_resnet_v2.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from models.inception_modules import Stem, ReductionA, BasicConv2D, Conv2DLinear 3 | from configuration import NUM_CLASSES 4 | 5 | 6 | class InceptionResNetA(tf.keras.layers.Layer): 7 | def __init__(self): 8 | super(InceptionResNetA, self).__init__() 9 | self.b1_conv = BasicConv2D(filters=32, 10 | kernel_size=(1, 1), 11 | strides=1, 12 | padding="same") 13 | self.b2_conv1 = BasicConv2D(filters=32, 14 | kernel_size=(1, 1), 15 | strides=1, 16 | padding="same") 17 | self.b2_conv2 = BasicConv2D(filters=32, 18 | kernel_size=(3, 3), 19 | strides=1, 20 | padding="same") 21 | self.b3_conv1 = BasicConv2D(filters=32, 22 | kernel_size=(1, 1), 23 | strides=1, 24 | padding="same") 25 | self.b3_conv2 = BasicConv2D(filters=48, 26 | kernel_size=(3, 3), 27 | strides=1, 28 | padding="same") 29 | self.b3_conv3 = BasicConv2D(filters=64, 30 | kernel_size=(3, 3), 31 | strides=1, 32 | padding="same") 33 | self.conv = Conv2DLinear(filters=384, 34 | kernel_size=(1, 1), 35 | strides=1, 36 | padding="same") 37 | 38 | def call(self, inputs, training=None, **kwargs): 39 | b1 = self.b1_conv(inputs, training=training) 40 | b2 = self.b2_conv1(inputs, training=training) 41 | b2 = self.b2_conv2(b2, training=training) 42 | b3 = self.b3_conv1(inputs, training=training) 43 | b3 = self.b3_conv2(b3, training=training) 44 | b3 = self.b3_conv3(b3, training=training) 45 | 46 | x = tf.concat(values=[b1, b2, b3], axis=-1) 47 | x = self.conv(x, training=training) 48 | 49 | output = tf.keras.layers.add([x, inputs]) 50 | return tf.nn.relu(output) 51 | 52 | 53 | class InceptionResNetB(tf.keras.layers.Layer): 54 | def __init__(self): 55 | super(InceptionResNetB, self).__init__() 56 | self.b1_conv = BasicConv2D(filters=192, 57 | kernel_size=(1, 1), 58 | strides=1, 59 | padding="same") 60 | self.b2_conv1 = BasicConv2D(filters=128, 61 | kernel_size=(1, 1), 62 | strides=1, 63 | padding="same") 64 | self.b2_conv2 = BasicConv2D(filters=160, 65 | kernel_size=(1, 7), 66 | strides=1, 67 | padding="same") 68 | self.b2_conv3 = BasicConv2D(filters=192, 69 | kernel_size=(7, 1), 70 | strides=1, 71 | padding="same") 72 | self.conv = Conv2DLinear(filters=1152, 73 | kernel_size=(1, 1), 74 | strides=1, 75 | padding="same") 76 | 77 | def call(self, inputs, training=None, **kwargs): 78 | b1 = self.b1_conv(inputs, training=training) 79 | b2 = self.b2_conv1(inputs, training=training) 80 | b2 = self.b2_conv2(b2, training=training) 81 | b2 = self.b2_conv3(b2, training=training) 82 | 83 | x = tf.concat(values=[b1, b2], axis=-1) 84 | x = self.conv(x, training=training) 85 | 86 | output = tf.keras.layers.add([x, inputs]) 87 | 88 | return tf.nn.relu(output) 89 | 90 | 91 | class InceptionResNetC(tf.keras.layers.Layer): 92 | def __init__(self): 93 | super(InceptionResNetC, self).__init__() 94 | self.b1_conv = BasicConv2D(filters=192, 95 | kernel_size=(1, 1), 96 | strides=1, 97 | padding="same") 98 | self.b2_conv1 = BasicConv2D(filters=192, 99 | kernel_size=(1, 1), 100 | strides=1, 101 | padding="same") 102 | self.b2_conv2 = BasicConv2D(filters=224, 103 | kernel_size=(1, 3), 104 | strides=1, 105 | padding="same") 106 | self.b2_conv3 = BasicConv2D(filters=256, 107 | kernel_size=(3, 1), 108 | strides=1, 109 | padding="same") 110 | self.conv = Conv2DLinear(filters=2144, 111 | kernel_size=(1, 1), 112 | strides=1, 113 | padding="same") 114 | 115 | def call(self, inputs, training=None, **kwargs): 116 | b1 = self.b1_conv(inputs, training=training) 117 | b2 = self.b2_conv1(inputs, training=training) 118 | b2 = self.b2_conv2(b2, training=training) 119 | b2 = self.b2_conv3(b2, training=training) 120 | 121 | x = tf.concat(values=[b1, b2], axis=-1) 122 | x = self.conv(x, training=training) 123 | 124 | output = tf.keras.layers.add([x, inputs]) 125 | 126 | return tf.nn.relu(output) 127 | 128 | 129 | class ReductionB(tf.keras.layers.Layer): 130 | def __init__(self): 131 | super(ReductionB, self).__init__() 132 | self.b1_maxpool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 133 | strides=2, 134 | padding="valid") 135 | self.b2_conv1 = BasicConv2D(filters=256, 136 | kernel_size=(1, 1), 137 | strides=1, 138 | padding="same") 139 | self.b2_conv2 = BasicConv2D(filters=384, 140 | kernel_size=(3, 3), 141 | strides=2, 142 | padding="valid") 143 | self.b3_conv1 = BasicConv2D(filters=256, 144 | kernel_size=(1, 1), 145 | strides=1, 146 | padding="same") 147 | self.b3_conv2 = BasicConv2D(filters=288, 148 | kernel_size=(3, 3), 149 | strides=2, 150 | padding="valid") 151 | self.b4_conv1 = BasicConv2D(filters=256, 152 | kernel_size=(1, 1), 153 | strides=1, 154 | padding="same") 155 | self.b4_conv2 = BasicConv2D(filters=288, 156 | kernel_size=(3, 3), 157 | strides=1, 158 | padding="same") 159 | self.b4_conv3 = BasicConv2D(filters=320, 160 | kernel_size=(3, 3), 161 | strides=2, 162 | padding="valid") 163 | 164 | def call(self, inputs, training=None, **kwargs): 165 | b1 = self.b1_maxpool(inputs) 166 | 167 | b2 = self.b2_conv1(inputs, training=training) 168 | b2 = self.b2_conv2(b2, training=training) 169 | 170 | b3 = self.b3_conv1(inputs, training=training) 171 | b3 = self.b3_conv2(b3, training=training) 172 | 173 | b4 = self.b4_conv1(inputs, training=training) 174 | b4 = self.b4_conv2(b4, training=training) 175 | b4 = self.b4_conv3(b4, training=training) 176 | 177 | return tf.concat(values=[b1, b2, b3, b4], axis=-1) 178 | 179 | 180 | def build_inception_resnet_a(n): 181 | block = tf.keras.Sequential() 182 | for _ in range(n): 183 | block.add(InceptionResNetA()) 184 | return block 185 | 186 | 187 | def build_inception_resnet_b(n): 188 | block = tf.keras.Sequential() 189 | for _ in range(n): 190 | block.add(InceptionResNetB()) 191 | return block 192 | 193 | 194 | def build_inception_resnet_c(n): 195 | block = tf.keras.Sequential() 196 | for _ in range(n): 197 | block.add(InceptionResNetC()) 198 | return block 199 | 200 | 201 | class InceptionResNetV2(tf.keras.Model): 202 | def __init__(self): 203 | super(InceptionResNetV2, self).__init__() 204 | self.stem = Stem() 205 | self.inception_resnet_a = build_inception_resnet_a(5) 206 | self.reduction_a = ReductionA(k=256, l=256, m=384, n=384) 207 | self.inception_resnet_b = build_inception_resnet_b(10) 208 | self.reduction_b = ReductionB() 209 | self.inception_resnet_c = build_inception_resnet_c(5) 210 | self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(8, 8)) 211 | self.dropout = tf.keras.layers.Dropout(rate=0.2) 212 | self.flat = tf.keras.layers.Flatten() 213 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, 214 | activation=tf.keras.activations.softmax) 215 | 216 | def call(self, inputs, training=None, mask=None): 217 | x = self.stem(inputs, training=training) 218 | x = self.inception_resnet_a(x, training=training) 219 | x = self.reduction_a(x, training=training) 220 | x = self.inception_resnet_b(x, training=training) 221 | x = self.reduction_b(x, training=training) 222 | x = self.inception_resnet_c(x, training=training) 223 | x = self.avgpool(x) 224 | x = self.dropout(x, training=training) 225 | x = self.flat(x) 226 | x = self.fc(x) 227 | 228 | return x 229 | 230 | def __repr__(self): 231 | return "InceptionResNetV2" 232 | -------------------------------------------------------------------------------- /models/inception_v4.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from models.inception_modules import Stem, InceptionBlockA, InceptionBlockB, \ 3 | InceptionBlockC, ReductionA, ReductionB 4 | from configuration import NUM_CLASSES 5 | 6 | 7 | def build_inception_block_a(n): 8 | block = tf.keras.Sequential() 9 | for _ in range(n): 10 | block.add(InceptionBlockA()) 11 | return block 12 | 13 | 14 | def build_inception_block_b(n): 15 | block = tf.keras.Sequential() 16 | for _ in range(n): 17 | block.add(InceptionBlockB()) 18 | return block 19 | 20 | 21 | def build_inception_block_c(n): 22 | block = tf.keras.Sequential() 23 | for _ in range(n): 24 | block.add(InceptionBlockC()) 25 | return block 26 | 27 | 28 | class InceptionV4(tf.keras.Model): 29 | def __init__(self): 30 | super(InceptionV4, self).__init__() 31 | self.stem = Stem() 32 | self.inception_a = build_inception_block_a(4) 33 | self.reduction_a = ReductionA(k=192, l=224, m=256, n=384) 34 | self.inception_b = build_inception_block_b(7) 35 | self.reduction_b = ReductionB() 36 | self.inception_c = build_inception_block_c(3) 37 | self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(8, 8)) 38 | self.dropout = tf.keras.layers.Dropout(rate=0.2) 39 | self.flat = tf.keras.layers.Flatten() 40 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, 41 | activation=tf.keras.activations.softmax) 42 | 43 | def call(self, inputs, training=True, mask=None): 44 | x = self.stem(inputs, training=training) 45 | x = self.inception_a(x, training=training) 46 | x = self.reduction_a(x, training=training) 47 | x = self.inception_b(x, training=training) 48 | x = self.reduction_b(x, training=training) 49 | x = self.inception_c(x, training=training) 50 | x = self.avgpool(x) 51 | x = self.dropout(x, training=training) 52 | x = self.flat(x) 53 | x = self.fc(x) 54 | 55 | return x 56 | 57 | def __repr__(self): 58 | return "InceptionV4" 59 | -------------------------------------------------------------------------------- /models/mobilenet_v1.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from configuration import NUM_CLASSES 3 | 4 | import tensorflow.keras as nn 5 | 6 | 7 | class SeparableConv2D(nn.layers.Layer): 8 | def __init__(self, filters, kernel_size, strides=1, padding="same"): 9 | super(SeparableConv2D, self).__init__() 10 | self.x1 = nn.layers.SeparableConv2D(filters=filters, 11 | kernel_size=kernel_size, 12 | strides=strides, 13 | padding=padding) 14 | self.x2 = nn.layers.BatchNormalization() 15 | self.x3 = nn.layers.ReLU() 16 | 17 | def call(self, inputs, training=None, *args, **kwargs): 18 | x = self.x1(inputs) 19 | x = self.x2(x, training=training) 20 | x = self.x3(x) 21 | return x 22 | 23 | 24 | class MobileNetV1(tf.keras.Model): 25 | def __init__(self): 26 | super(MobileNetV1, self).__init__() 27 | self.conv1 = tf.keras.layers.Conv2D(filters=32, 28 | kernel_size=(3, 3), 29 | strides=2, 30 | padding="same") 31 | self.separable_conv_1 = SeparableConv2D(filters=64, 32 | kernel_size=(3, 3), 33 | strides=1, 34 | padding="same") 35 | self.separable_conv_2 = SeparableConv2D(filters=128, 36 | kernel_size=(3, 3), 37 | strides=2, 38 | padding="same") 39 | self.separable_conv_3 = SeparableConv2D(filters=128, 40 | kernel_size=(3, 3), 41 | strides=1, 42 | padding="same") 43 | self.separable_conv_4 = SeparableConv2D(filters=256, 44 | kernel_size=(3, 3), 45 | strides=2, 46 | padding="same") 47 | self.separable_conv_5 = SeparableConv2D(filters=256, 48 | kernel_size=(3, 3), 49 | strides=1, 50 | padding="same") 51 | self.separable_conv_6 = SeparableConv2D(filters=512, 52 | kernel_size=(3, 3), 53 | strides=2, 54 | padding="same") 55 | 56 | self.separable_conv_7 = SeparableConv2D(filters=512, 57 | kernel_size=(3, 3), 58 | strides=1, 59 | padding="same") 60 | self.separable_conv_8 = SeparableConv2D(filters=512, 61 | kernel_size=(3, 3), 62 | strides=1, 63 | padding="same") 64 | self.separable_conv_9 = SeparableConv2D(filters=512, 65 | kernel_size=(3, 3), 66 | strides=1, 67 | padding="same") 68 | self.separable_conv_10 = SeparableConv2D(filters=512, 69 | kernel_size=(3, 3), 70 | strides=1, 71 | padding="same") 72 | self.separable_conv_11 = SeparableConv2D(filters=512, 73 | kernel_size=(3, 3), 74 | strides=1, 75 | padding="same") 76 | 77 | self.separable_conv_12 = SeparableConv2D(filters=1024, 78 | kernel_size=(3, 3), 79 | strides=2, 80 | padding="same") 81 | self.separable_conv_13 = SeparableConv2D(filters=1024, 82 | kernel_size=(3, 3), 83 | strides=1, 84 | padding="same") 85 | 86 | self.avg_pool = tf.keras.layers.AveragePooling2D(pool_size=(7, 7), 87 | strides=1) 88 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, 89 | activation=tf.keras.activations.softmax) 90 | 91 | def call(self, inputs, training=None, mask=None): 92 | x = self.conv1(inputs) 93 | x = self.separable_conv_1(x, training=training) 94 | x = self.separable_conv_2(x, training=training) 95 | x = self.separable_conv_3(x, training=training) 96 | x = self.separable_conv_4(x, training=training) 97 | x = self.separable_conv_5(x, training=training) 98 | x = self.separable_conv_6(x, training=training) 99 | x = self.separable_conv_7(x, training=training) 100 | x = self.separable_conv_8(x, training=training) 101 | x = self.separable_conv_9(x, training=training) 102 | x = self.separable_conv_10(x, training=training) 103 | x = self.separable_conv_11(x, training=training) 104 | x = self.separable_conv_12(x, training=training) 105 | x = self.separable_conv_13(x, training=training) 106 | 107 | x = self.avg_pool(x) 108 | x = self.fc(x) 109 | 110 | return x 111 | 112 | def __repr__(self): 113 | return "MobileNetV1" 114 | -------------------------------------------------------------------------------- /models/mobilenet_v2.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from configuration import NUM_CLASSES 3 | 4 | 5 | class BottleNeck(tf.keras.layers.Layer): 6 | def __init__(self, input_channels, output_channels, expansion_factor, stride): 7 | self.stride = stride 8 | self.input_channels = input_channels 9 | self.output_channels = output_channels 10 | super(BottleNeck, self).__init__() 11 | self.conv1 = tf.keras.layers.Conv2D(filters=input_channels * expansion_factor, 12 | kernel_size=(1, 1), 13 | strides=1, 14 | padding="same") 15 | self.bn1 = tf.keras.layers.BatchNormalization() 16 | self.dwconv = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3), 17 | strides=stride, 18 | padding="same") 19 | self.bn2 = tf.keras.layers.BatchNormalization() 20 | self.conv2 = tf.keras.layers.Conv2D(filters=output_channels, 21 | kernel_size=(1, 1), 22 | strides=1, 23 | padding="same") 24 | self.bn3 = tf.keras.layers.BatchNormalization() 25 | self.linear = tf.keras.layers.Activation(tf.keras.activations.linear) 26 | 27 | def call(self, inputs, training=None, **kwargs): 28 | x = self.conv1(inputs) 29 | x = self.bn1(x, training=training) 30 | x = tf.nn.relu6(x) 31 | x = self.dwconv(x) 32 | x = self.bn2(x, training=training) 33 | x = tf.nn.relu6(x) 34 | x = self.conv2(x) 35 | x = self.bn3(x, training=training) 36 | x = self.linear(x) 37 | if self.stride == 1 and self.input_channels == self.output_channels: 38 | x = tf.keras.layers.add([x, inputs]) 39 | return x 40 | 41 | 42 | def build_bottleneck(t, in_channel_num, out_channel_num, n, s): 43 | # t : expansion factor 44 | # s : stride 45 | # n : repeat times 46 | bottleneck = tf.keras.Sequential() 47 | for i in range(n): 48 | if i == 0: 49 | bottleneck.add(BottleNeck(input_channels=in_channel_num, 50 | output_channels=out_channel_num, 51 | expansion_factor=t, 52 | stride=s)) 53 | else: 54 | bottleneck.add(BottleNeck(input_channels=out_channel_num, 55 | output_channels=out_channel_num, 56 | expansion_factor=t, 57 | stride=1)) 58 | 59 | return bottleneck 60 | 61 | 62 | class MobileNetV2(tf.keras.Model): 63 | def __init__(self): 64 | super(MobileNetV2, self).__init__() 65 | self.conv1 = tf.keras.layers.Conv2D(filters=32, 66 | kernel_size=(3, 3), 67 | strides=2, 68 | padding="same") 69 | self.bottleneck_1 = build_bottleneck(t=1, 70 | in_channel_num=32, 71 | out_channel_num=16, 72 | n=1, 73 | s=1) 74 | self.bottleneck_2 = build_bottleneck(t=6, 75 | in_channel_num=16, 76 | out_channel_num=24, 77 | n=2, 78 | s=2) 79 | self.bottleneck_3 = build_bottleneck(t=6, 80 | in_channel_num=24, 81 | out_channel_num=32, 82 | n=3, 83 | s=2) 84 | self.bottleneck_4 = build_bottleneck(t=6, 85 | in_channel_num=32, 86 | out_channel_num=64, 87 | n=4, 88 | s=2) 89 | self.bottleneck_5 = build_bottleneck(t=6, 90 | in_channel_num=64, 91 | out_channel_num=96, 92 | n=3, 93 | s=1) 94 | self.bottleneck_6 = build_bottleneck(t=6, 95 | in_channel_num=96, 96 | out_channel_num=160, 97 | n=3, 98 | s=2) 99 | self.bottleneck_7 = build_bottleneck(t=6, 100 | in_channel_num=160, 101 | out_channel_num=320, 102 | n=1, 103 | s=1) 104 | 105 | self.conv2 = tf.keras.layers.Conv2D(filters=1280, 106 | kernel_size=(1, 1), 107 | strides=1, 108 | padding="same") 109 | self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(7, 7)) 110 | self.conv3 = tf.keras.layers.Conv2D(filters=NUM_CLASSES, 111 | kernel_size=(1, 1), 112 | strides=1, 113 | padding="same", 114 | activation=tf.keras.activations.softmax) 115 | 116 | def call(self, inputs, training=None, mask=None): 117 | x = self.conv1(inputs) 118 | x = self.bottleneck_1(x, training=training) 119 | x = self.bottleneck_2(x, training=training) 120 | x = self.bottleneck_3(x, training=training) 121 | x = self.bottleneck_4(x, training=training) 122 | x = self.bottleneck_5(x, training=training) 123 | x = self.bottleneck_6(x, training=training) 124 | x = self.bottleneck_7(x, training=training) 125 | 126 | x = self.conv2(x) 127 | x = self.avgpool(x) 128 | x = self.conv3(x) 129 | 130 | return x 131 | 132 | def __repr__(self): 133 | return "MobileNetV2" 134 | 135 | -------------------------------------------------------------------------------- /models/mobilenet_v3_block.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | 4 | def h_sigmoid(x): 5 | return tf.nn.relu6(x + 3) / 6 6 | 7 | 8 | def h_swish(x): 9 | return x * h_sigmoid(x) 10 | 11 | 12 | class SEBlock(tf.keras.layers.Layer): 13 | def __init__(self, input_channels, r=16): 14 | super(SEBlock, self).__init__() 15 | self.pool = tf.keras.layers.GlobalAveragePooling2D() 16 | self.fc1 = tf.keras.layers.Dense(units=input_channels // r) 17 | self.fc2 = tf.keras.layers.Dense(units=input_channels) 18 | 19 | def call(self, inputs, **kwargs): 20 | branch = self.pool(inputs) 21 | branch = self.fc1(branch) 22 | branch = tf.nn.relu(branch) 23 | branch = self.fc2(branch) 24 | branch = h_sigmoid(branch) 25 | branch = tf.expand_dims(input=branch, axis=1) 26 | branch = tf.expand_dims(input=branch, axis=1) 27 | output = inputs * branch 28 | return output 29 | 30 | 31 | class BottleNeck(tf.keras.layers.Layer): 32 | def __init__(self, in_size, exp_size, out_size, s, is_se_existing, NL, k): 33 | super(BottleNeck, self).__init__() 34 | self.stride = s 35 | self.in_size = in_size 36 | self.out_size = out_size 37 | self.is_se_existing = is_se_existing 38 | self.NL = NL 39 | self.conv1 = tf.keras.layers.Conv2D(filters=exp_size, 40 | kernel_size=(1, 1), 41 | strides=1, 42 | padding="same") 43 | self.bn1 = tf.keras.layers.BatchNormalization() 44 | self.dwconv = tf.keras.layers.DepthwiseConv2D(kernel_size=(k, k), 45 | strides=s, 46 | padding="same") 47 | self.bn2 = tf.keras.layers.BatchNormalization() 48 | self.se = SEBlock(input_channels=exp_size) 49 | self.conv2 = tf.keras.layers.Conv2D(filters=out_size, 50 | kernel_size=(1, 1), 51 | strides=1, 52 | padding="same") 53 | self.bn3 = tf.keras.layers.BatchNormalization() 54 | self.linear = tf.keras.layers.Activation(tf.keras.activations.linear) 55 | 56 | def call(self, inputs, training=None, **kwargs): 57 | x = self.conv1(inputs) 58 | x = self.bn1(x, training=training) 59 | if self.NL == "HS": 60 | x = h_swish(x) 61 | elif self.NL == "RE": 62 | x = tf.nn.relu6(x) 63 | x = self.dwconv(x) 64 | x = self.bn2(x, training=training) 65 | if self.NL == "HS": 66 | x = h_swish(x) 67 | elif self.NL == "RE": 68 | x = tf.nn.relu6(x) 69 | if self.is_se_existing: 70 | x = self.se(x) 71 | x = self.conv2(x) 72 | x = self.bn3(x, training=training) 73 | x = self.linear(x) 74 | 75 | if self.stride == 1 and self.in_size == self.out_size: 76 | x = tf.keras.layers.add([x, inputs]) 77 | 78 | return x 79 | -------------------------------------------------------------------------------- /models/mobilenet_v3_large.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from models.mobilenet_v3_block import BottleNeck, h_swish 3 | from configuration import NUM_CLASSES 4 | 5 | 6 | class MobileNetV3Large(tf.keras.Model): 7 | def __init__(self): 8 | super(MobileNetV3Large, self).__init__() 9 | self.conv1 = tf.keras.layers.Conv2D(filters=16, 10 | kernel_size=(3, 3), 11 | strides=2, 12 | padding="same") 13 | self.bn1 = tf.keras.layers.BatchNormalization() 14 | self.bneck1 = BottleNeck(in_size=16, exp_size=16, out_size=16, s=1, is_se_existing=False, NL="RE", k=3) 15 | self.bneck2 = BottleNeck(in_size=16, exp_size=64, out_size=24, s=2, is_se_existing=False, NL="RE", k=3) 16 | self.bneck3 = BottleNeck(in_size=24, exp_size=72, out_size=24, s=1, is_se_existing=False, NL="RE", k=3) 17 | self.bneck4 = BottleNeck(in_size=24, exp_size=72, out_size=40, s=2, is_se_existing=True, NL="RE", k=5) 18 | self.bneck5 = BottleNeck(in_size=40, exp_size=120, out_size=40, s=1, is_se_existing=True, NL="RE", k=5) 19 | self.bneck6 = BottleNeck(in_size=40, exp_size=120, out_size=40, s=1, is_se_existing=True, NL="RE", k=5) 20 | self.bneck7 = BottleNeck(in_size=40, exp_size=240, out_size=80, s=2, is_se_existing=False, NL="HS", k=3) 21 | self.bneck8 = BottleNeck(in_size=80, exp_size=200, out_size=80, s=1, is_se_existing=False, NL="HS", k=3) 22 | self.bneck9 = BottleNeck(in_size=80, exp_size=184, out_size=80, s=1, is_se_existing=False, NL="HS", k=3) 23 | self.bneck10 = BottleNeck(in_size=80, exp_size=184, out_size=80, s=1, is_se_existing=False, NL="HS", k=3) 24 | self.bneck11 = BottleNeck(in_size=80, exp_size=480, out_size=112, s=1, is_se_existing=True, NL="HS", k=3) 25 | self.bneck12 = BottleNeck(in_size=112, exp_size=672, out_size=112, s=1, is_se_existing=True, NL="HS", k=3) 26 | self.bneck13 = BottleNeck(in_size=112, exp_size=672, out_size=160, s=2, is_se_existing=True, NL="HS", k=5) 27 | self.bneck14 = BottleNeck(in_size=160, exp_size=960, out_size=160, s=1, is_se_existing=True, NL="HS", k=5) 28 | self.bneck15 = BottleNeck(in_size=160, exp_size=960, out_size=160, s=1, is_se_existing=True, NL="HS", k=5) 29 | 30 | self.conv2 = tf.keras.layers.Conv2D(filters=960, 31 | kernel_size=(1, 1), 32 | strides=1, 33 | padding="same") 34 | self.bn2 = tf.keras.layers.BatchNormalization() 35 | self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(7, 7), 36 | strides=1) 37 | self.conv3 = tf.keras.layers.Conv2D(filters=1280, 38 | kernel_size=(1, 1), 39 | strides=1, 40 | padding="same") 41 | self.conv4 = tf.keras.layers.Conv2D(filters=NUM_CLASSES, 42 | kernel_size=(1, 1), 43 | strides=1, 44 | padding="same", 45 | activation=tf.keras.activations.softmax) 46 | 47 | def call(self, inputs, training=None, mask=None): 48 | x = self.conv1(inputs) 49 | x = self.bn1(x, training=training) 50 | x = h_swish(x) 51 | 52 | x = self.bneck1(x, training=training) 53 | x = self.bneck2(x, training=training) 54 | x = self.bneck3(x, training=training) 55 | x = self.bneck4(x, training=training) 56 | x = self.bneck5(x, training=training) 57 | x = self.bneck6(x, training=training) 58 | x = self.bneck7(x, training=training) 59 | x = self.bneck8(x, training=training) 60 | x = self.bneck9(x, training=training) 61 | x = self.bneck10(x, training=training) 62 | x = self.bneck11(x, training=training) 63 | x = self.bneck12(x, training=training) 64 | x = self.bneck13(x, training=training) 65 | x = self.bneck14(x, training=training) 66 | x = self.bneck15(x, training=training) 67 | 68 | x = self.conv2(x) 69 | x = self.bn2(x, training=training) 70 | x = h_swish(x) 71 | x = self.avgpool(x) 72 | x = self.conv3(x) 73 | x = h_swish(x) 74 | x = self.conv4(x) 75 | 76 | return x 77 | 78 | def __repr__(self): 79 | return "MobileNetV3Large" 80 | -------------------------------------------------------------------------------- /models/mobilenet_v3_small.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from models.mobilenet_v3_block import BottleNeck, h_swish 3 | from configuration import NUM_CLASSES 4 | 5 | 6 | class MobileNetV3Small(tf.keras.Model): 7 | def __init__(self): 8 | super(MobileNetV3Small, self).__init__() 9 | self.conv1 = tf.keras.layers.Conv2D(filters=16, 10 | kernel_size=(3, 3), 11 | strides=2, 12 | padding="same") 13 | self.bn1 = tf.keras.layers.BatchNormalization() 14 | self.bneck1 = BottleNeck(in_size=16, exp_size=16, out_size=16, s=2, is_se_existing=True, NL="RE", k=3) 15 | self.bneck2 = BottleNeck(in_size=16, exp_size=72, out_size=24, s=2, is_se_existing=False, NL="RE", k=3) 16 | self.bneck3 = BottleNeck(in_size=24, exp_size=88, out_size=24, s=1, is_se_existing=False, NL="RE", k=3) 17 | self.bneck4 = BottleNeck(in_size=24, exp_size=96, out_size=40, s=2, is_se_existing=True, NL="HS", k=5) 18 | self.bneck5 = BottleNeck(in_size=40, exp_size=240, out_size=40, s=1, is_se_existing=True, NL="HS", k=5) 19 | self.bneck6 = BottleNeck(in_size=40, exp_size=240, out_size=40, s=1, is_se_existing=True, NL="HS", k=5) 20 | self.bneck7 = BottleNeck(in_size=40, exp_size=120, out_size=48, s=1, is_se_existing=True, NL="HS", k=5) 21 | self.bneck8 = BottleNeck(in_size=48, exp_size=144, out_size=48, s=1, is_se_existing=True, NL="HS", k=5) 22 | self.bneck9 = BottleNeck(in_size=48, exp_size=288, out_size=96, s=2, is_se_existing=True, NL="HS", k=5) 23 | self.bneck10 = BottleNeck(in_size=96, exp_size=576, out_size=96, s=1, is_se_existing=True, NL="HS", k=5) 24 | self.bneck11 = BottleNeck(in_size=96, exp_size=576, out_size=96, s=1, is_se_existing=True, NL="HS", k=5) 25 | 26 | self.conv2 = tf.keras.layers.Conv2D(filters=576, 27 | kernel_size=(1, 1), 28 | strides=1, 29 | padding="same") 30 | self.bn2 = tf.keras.layers.BatchNormalization() 31 | self.avgpool = tf.keras.layers.AveragePooling2D(pool_size=(7, 7), 32 | strides=1) 33 | self.conv3 = tf.keras.layers.Conv2D(filters=1280, 34 | kernel_size=(1, 1), 35 | strides=1, 36 | padding="same") 37 | self.conv4 = tf.keras.layers.Conv2D(filters=NUM_CLASSES, 38 | kernel_size=(1, 1), 39 | strides=1, 40 | padding="same", 41 | activation=tf.keras.activations.softmax) 42 | 43 | def call(self, inputs, training=None, mask=None): 44 | x = self.conv1(inputs) 45 | x = self.bn1(x, training=training) 46 | x = h_swish(x) 47 | 48 | x = self.bneck1(x, training=training) 49 | x = self.bneck2(x, training=training) 50 | x = self.bneck3(x, training=training) 51 | x = self.bneck4(x, training=training) 52 | x = self.bneck5(x, training=training) 53 | x = self.bneck6(x, training=training) 54 | x = self.bneck7(x, training=training) 55 | x = self.bneck8(x, training=training) 56 | x = self.bneck9(x, training=training) 57 | x = self.bneck10(x, training=training) 58 | x = self.bneck11(x, training=training) 59 | 60 | x = self.conv2(x) 61 | x = self.bn2(x, training=training) 62 | x = h_swish(x) 63 | x = self.avgpool(x) 64 | x = self.conv3(x) 65 | x = h_swish(x) 66 | x = self.conv4(x) 67 | 68 | return x 69 | 70 | def __repr__(self): 71 | return "MobileNetV3Small" 72 | -------------------------------------------------------------------------------- /models/residual_block.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | 4 | class BasicBlock(tf.keras.layers.Layer): 5 | 6 | def __init__(self, filter_num, stride=1): 7 | super(BasicBlock, self).__init__() 8 | self.conv1 = tf.keras.layers.Conv2D(filters=filter_num, 9 | kernel_size=(3, 3), 10 | strides=stride, 11 | padding="same") 12 | self.bn1 = tf.keras.layers.BatchNormalization() 13 | self.conv2 = tf.keras.layers.Conv2D(filters=filter_num, 14 | kernel_size=(3, 3), 15 | strides=1, 16 | padding="same") 17 | self.bn2 = tf.keras.layers.BatchNormalization() 18 | if stride != 1: 19 | self.downsample = tf.keras.Sequential() 20 | self.downsample.add(tf.keras.layers.Conv2D(filters=filter_num, 21 | kernel_size=(1, 1), 22 | strides=stride)) 23 | self.downsample.add(tf.keras.layers.BatchNormalization()) 24 | else: 25 | self.downsample = lambda x: x 26 | 27 | def call(self, inputs, training=None, **kwargs): 28 | residual = self.downsample(inputs) 29 | 30 | x = self.conv1(inputs) 31 | x = self.bn1(x, training=training) 32 | x = tf.nn.relu(x) 33 | x = self.conv2(x) 34 | x = self.bn2(x, training=training) 35 | 36 | output = tf.nn.relu(tf.keras.layers.add([residual, x])) 37 | 38 | return output 39 | 40 | 41 | class BottleNeck(tf.keras.layers.Layer): 42 | def __init__(self, filter_num, stride=1): 43 | super(BottleNeck, self).__init__() 44 | self.conv1 = tf.keras.layers.Conv2D(filters=filter_num, 45 | kernel_size=(1, 1), 46 | strides=1, 47 | padding='same') 48 | self.bn1 = tf.keras.layers.BatchNormalization() 49 | self.conv2 = tf.keras.layers.Conv2D(filters=filter_num, 50 | kernel_size=(3, 3), 51 | strides=stride, 52 | padding='same') 53 | self.bn2 = tf.keras.layers.BatchNormalization() 54 | self.conv3 = tf.keras.layers.Conv2D(filters=filter_num * 4, 55 | kernel_size=(1, 1), 56 | strides=1, 57 | padding='same') 58 | self.bn3 = tf.keras.layers.BatchNormalization() 59 | 60 | self.downsample = tf.keras.Sequential() 61 | self.downsample.add(tf.keras.layers.Conv2D(filters=filter_num * 4, 62 | kernel_size=(1, 1), 63 | strides=stride)) 64 | self.downsample.add(tf.keras.layers.BatchNormalization()) 65 | 66 | def call(self, inputs, training=None, **kwargs): 67 | residual = self.downsample(inputs) 68 | 69 | x = self.conv1(inputs) 70 | x = self.bn1(x, training=training) 71 | x = tf.nn.relu(x) 72 | x = self.conv2(x) 73 | x = self.bn2(x, training=training) 74 | x = tf.nn.relu(x) 75 | x = self.conv3(x) 76 | x = self.bn3(x, training=training) 77 | 78 | output = tf.nn.relu(tf.keras.layers.add([residual, x])) 79 | 80 | return output 81 | 82 | 83 | def make_basic_block_layer(filter_num, blocks, stride=1): 84 | res_block = tf.keras.Sequential() 85 | res_block.add(BasicBlock(filter_num, stride=stride)) 86 | 87 | for _ in range(1, blocks): 88 | res_block.add(BasicBlock(filter_num, stride=1)) 89 | 90 | return res_block 91 | 92 | 93 | def make_bottleneck_layer(filter_num, blocks, stride=1): 94 | res_block = tf.keras.Sequential() 95 | res_block.add(BottleNeck(filter_num, stride=stride)) 96 | 97 | for _ in range(1, blocks): 98 | res_block.add(BottleNeck(filter_num, stride=1)) 99 | 100 | return res_block 101 | -------------------------------------------------------------------------------- /models/resnet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from configuration import NUM_CLASSES 3 | from models.residual_block import make_basic_block_layer, make_bottleneck_layer 4 | 5 | 6 | class ResNetTypeI(tf.keras.Model): 7 | def __init__(self, layer_params, model_name): 8 | super(ResNetTypeI, self).__init__() 9 | self.model_name = model_name 10 | 11 | self.conv1 = tf.keras.layers.Conv2D(filters=64, 12 | kernel_size=(7, 7), 13 | strides=2, 14 | padding="same") 15 | self.bn1 = tf.keras.layers.BatchNormalization() 16 | self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 17 | strides=2, 18 | padding="same") 19 | 20 | self.layer1 = make_basic_block_layer(filter_num=64, 21 | blocks=layer_params[0]) 22 | self.layer2 = make_basic_block_layer(filter_num=128, 23 | blocks=layer_params[1], 24 | stride=2) 25 | self.layer3 = make_basic_block_layer(filter_num=256, 26 | blocks=layer_params[2], 27 | stride=2) 28 | self.layer4 = make_basic_block_layer(filter_num=512, 29 | blocks=layer_params[3], 30 | stride=2) 31 | 32 | self.avgpool = tf.keras.layers.GlobalAveragePooling2D() 33 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, activation=tf.keras.activations.softmax) 34 | 35 | def call(self, inputs, training=None, mask=None): 36 | x = self.conv1(inputs) 37 | x = self.bn1(x, training=training) 38 | x = tf.nn.relu(x) 39 | x = self.pool1(x) 40 | x = self.layer1(x, training=training) 41 | x = self.layer2(x, training=training) 42 | x = self.layer3(x, training=training) 43 | x = self.layer4(x, training=training) 44 | x = self.avgpool(x) 45 | output = self.fc(x) 46 | 47 | return output 48 | 49 | def __repr__(self): 50 | return "ResNet_{}".format(self.model_name) 51 | 52 | 53 | class ResNetTypeII(tf.keras.Model): 54 | def __init__(self, layer_params, model_name): 55 | super(ResNetTypeII, self).__init__() 56 | self.model_name = model_name 57 | 58 | self.conv1 = tf.keras.layers.Conv2D(filters=64, 59 | kernel_size=(7, 7), 60 | strides=2, 61 | padding="same") 62 | self.bn1 = tf.keras.layers.BatchNormalization() 63 | self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 64 | strides=2, 65 | padding="same") 66 | 67 | self.layer1 = make_bottleneck_layer(filter_num=64, 68 | blocks=layer_params[0]) 69 | self.layer2 = make_bottleneck_layer(filter_num=128, 70 | blocks=layer_params[1], 71 | stride=2) 72 | self.layer3 = make_bottleneck_layer(filter_num=256, 73 | blocks=layer_params[2], 74 | stride=2) 75 | self.layer4 = make_bottleneck_layer(filter_num=512, 76 | blocks=layer_params[3], 77 | stride=2) 78 | 79 | self.avgpool = tf.keras.layers.GlobalAveragePooling2D() 80 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, activation=tf.keras.activations.softmax) 81 | 82 | def call(self, inputs, training=None, mask=None): 83 | x = self.conv1(inputs) 84 | x = self.bn1(x, training=training) 85 | x = tf.nn.relu(x) 86 | x = self.pool1(x) 87 | x = self.layer1(x, training=training) 88 | x = self.layer2(x, training=training) 89 | x = self.layer3(x, training=training) 90 | x = self.layer4(x, training=training) 91 | x = self.avgpool(x) 92 | output = self.fc(x) 93 | 94 | return output 95 | 96 | def __repr__(self): 97 | return "ResNet_{}".format(self.model_name) 98 | 99 | 100 | def resnet_18(): 101 | return ResNetTypeI(layer_params=[2, 2, 2, 2], model_name="18") 102 | 103 | 104 | def resnet_34(): 105 | return ResNetTypeI(layer_params=[3, 4, 6, 3], model_name="34") 106 | 107 | 108 | def resnet_50(): 109 | return ResNetTypeII(layer_params=[3, 4, 6, 3], model_name="50") 110 | 111 | 112 | def resnet_101(): 113 | return ResNetTypeII(layer_params=[3, 4, 23, 3], model_name="101") 114 | 115 | 116 | def resnet_152(): 117 | return ResNetTypeII(layer_params=[3, 8, 36, 3], model_name="152") 118 | -------------------------------------------------------------------------------- /models/resnext.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from models.resnext_block import build_ResNeXt_block 3 | from configuration import NUM_CLASSES 4 | 5 | 6 | class ResNeXt(tf.keras.Model): 7 | def __init__(self, repeat_num_list, cardinality, model_name): 8 | if len(repeat_num_list) != 4: 9 | raise ValueError("The length of repeat_num_list must be four.") 10 | super(ResNeXt, self).__init__() 11 | self.model_name = model_name 12 | 13 | self.conv1 = tf.keras.layers.Conv2D(filters=64, 14 | kernel_size=(7, 7), 15 | strides=2, 16 | padding="same") 17 | self.bn1 = tf.keras.layers.BatchNormalization() 18 | self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 19 | strides=2, 20 | padding="same") 21 | self.block1 = build_ResNeXt_block(filters=128, 22 | strides=1, 23 | groups=cardinality, 24 | repeat_num=repeat_num_list[0]) 25 | self.block2 = build_ResNeXt_block(filters=256, 26 | strides=2, 27 | groups=cardinality, 28 | repeat_num=repeat_num_list[1]) 29 | self.block3 = build_ResNeXt_block(filters=512, 30 | strides=2, 31 | groups=cardinality, 32 | repeat_num=repeat_num_list[2]) 33 | self.block4 = build_ResNeXt_block(filters=1024, 34 | strides=2, 35 | groups=cardinality, 36 | repeat_num=repeat_num_list[3]) 37 | self.pool2 = tf.keras.layers.GlobalAveragePooling2D() 38 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, 39 | activation=tf.keras.activations.softmax) 40 | 41 | def call(self, inputs, training=None, mask=None): 42 | x = self.conv1(inputs) 43 | x = self.bn1(x, training=training) 44 | x = tf.nn.relu(x) 45 | x = self.pool1(x) 46 | 47 | x = self.block1(x, training=training) 48 | x = self.block2(x, training=training) 49 | x = self.block3(x, training=training) 50 | x = self.block4(x, training=training) 51 | 52 | x = self.pool2(x) 53 | x = self.fc(x) 54 | 55 | return x 56 | 57 | def __repr__(self): 58 | return "ResNeXt-{}".format(self.model_name) 59 | 60 | 61 | def ResNeXt50(): 62 | return ResNeXt(repeat_num_list=[3, 4, 6, 3], 63 | cardinality=32, 64 | model_name="50") 65 | 66 | 67 | def ResNeXt101(): 68 | return ResNeXt(repeat_num_list=[3, 4, 23, 3], 69 | cardinality=32, 70 | model_name="101") -------------------------------------------------------------------------------- /models/resnext_block.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from models.group_convolution import get_group_conv 4 | 5 | 6 | class ResNeXt_BottleNeck(tf.keras.layers.Layer): 7 | def __init__(self, filters, strides, groups): 8 | super(ResNeXt_BottleNeck, self).__init__() 9 | self.conv1 = tf.keras.layers.Conv2D(filters=filters, 10 | kernel_size=(1, 1), 11 | strides=1, 12 | padding="same") 13 | self.bn1 = tf.keras.layers.BatchNormalization() 14 | # self.group_conv = tf.keras.layers.Conv2D(filters=filters, 15 | # kernel_size=(3, 3), 16 | # strides=strides, 17 | # padding="same", 18 | # groups=groups) 19 | self.group_conv = get_group_conv(in_channels=filters, 20 | out_channels=filters, 21 | kernel_size=(3, 3), 22 | strides=strides, 23 | padding="same", 24 | groups=groups) 25 | self.bn2 = tf.keras.layers.BatchNormalization() 26 | self.conv2 = tf.keras.layers.Conv2D(filters=2 * filters, 27 | kernel_size=(1, 1), 28 | strides=1, 29 | padding="same") 30 | self.bn3 = tf.keras.layers.BatchNormalization() 31 | self.shortcut_conv = tf.keras.layers.Conv2D(filters=2 * filters, 32 | kernel_size=(1, 1), 33 | strides=strides, 34 | padding="same") 35 | self.shortcut_bn = tf.keras.layers.BatchNormalization() 36 | 37 | def call(self, inputs, training=None, **kwargs): 38 | x = self.conv1(inputs) 39 | x = self.bn1(x, training=training) 40 | x = tf.nn.relu(x) 41 | x = self.group_conv(x) 42 | x = self.bn2(x, training=training) 43 | x = tf.nn.relu(x) 44 | x = self.conv2(x) 45 | x = self.bn3(x, training=training) 46 | 47 | shortcut = self.shortcut_conv(inputs) 48 | shortcut = self.shortcut_bn(shortcut, training=training) 49 | 50 | output = tf.nn.relu(tf.keras.layers.add([x, shortcut])) 51 | return output 52 | 53 | 54 | def build_ResNeXt_block(filters, strides, groups, repeat_num): 55 | block = tf.keras.Sequential() 56 | block.add(ResNeXt_BottleNeck(filters=filters, 57 | strides=strides, 58 | groups=groups)) 59 | for _ in range(1, repeat_num): 60 | block.add(ResNeXt_BottleNeck(filters=filters, 61 | strides=1, 62 | groups=groups)) 63 | 64 | return block -------------------------------------------------------------------------------- /models/se_resnet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from configuration import NUM_CLASSES 3 | 4 | 5 | class SEBlock(tf.keras.layers.Layer): 6 | def __init__(self, input_channels, r=16): 7 | super(SEBlock, self).__init__() 8 | self.pool = tf.keras.layers.GlobalAveragePooling2D() 9 | self.fc1 = tf.keras.layers.Dense(units=input_channels // r) 10 | self.fc2 = tf.keras.layers.Dense(units=input_channels) 11 | 12 | def call(self, inputs, **kwargs): 13 | branch = self.pool(inputs) 14 | branch = self.fc1(branch) 15 | branch = tf.nn.relu(branch) 16 | branch = self.fc2(branch) 17 | branch = tf.nn.sigmoid(branch) 18 | branch = tf.expand_dims(input=branch, axis=1) 19 | branch = tf.expand_dims(input=branch, axis=1) 20 | output = tf.keras.layers.multiply(inputs=[inputs, branch]) 21 | return output 22 | 23 | 24 | class BottleNeck(tf.keras.layers.Layer): 25 | def __init__(self, filter_num, stride=1): 26 | super(BottleNeck, self).__init__() 27 | self.conv1 = tf.keras.layers.Conv2D(filters=filter_num, 28 | kernel_size=(1, 1), 29 | strides=1, 30 | padding='same') 31 | self.bn1 = tf.keras.layers.BatchNormalization() 32 | self.conv2 = tf.keras.layers.Conv2D(filters=filter_num, 33 | kernel_size=(3, 3), 34 | strides=stride, 35 | padding='same') 36 | self.bn2 = tf.keras.layers.BatchNormalization() 37 | self.conv3 = tf.keras.layers.Conv2D(filters=filter_num * 4, 38 | kernel_size=(1, 1), 39 | strides=1, 40 | padding='same') 41 | self.bn3 = tf.keras.layers.BatchNormalization() 42 | self.se = SEBlock(input_channels=filter_num * 4) 43 | 44 | self.downsample = tf.keras.Sequential() 45 | self.downsample.add(tf.keras.layers.Conv2D(filters=filter_num * 4, 46 | kernel_size=(1, 1), 47 | strides=stride)) 48 | self.downsample.add(tf.keras.layers.BatchNormalization()) 49 | 50 | def call(self, inputs, training=None): 51 | identity = self.downsample(inputs) 52 | 53 | x = self.conv1(inputs) 54 | x = self.bn1(x, training=training) 55 | x = tf.nn.relu(x) 56 | x = self.conv2(x) 57 | x = self.bn2(x, training=training) 58 | x = tf.nn.relu(x) 59 | x = self.conv3(x) 60 | x = self.bn3(x, training=training) 61 | x = self.se(x) 62 | output = tf.nn.relu(tf.keras.layers.add([identity, x])) 63 | return output 64 | 65 | 66 | class SEResNet(tf.keras.Model): 67 | def __init__(self, block_num, model_name): 68 | super(SEResNet, self).__init__() 69 | self.model_name = model_name 70 | 71 | self.pre1 = tf.keras.layers.Conv2D(filters=64, 72 | kernel_size=(7, 7), 73 | strides=2, 74 | padding='same') 75 | self.pre2 = tf.keras.layers.BatchNormalization() 76 | self.pre3 = tf.keras.layers.Activation(tf.keras.activations.relu) 77 | self.pre4 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 78 | strides=2) 79 | 80 | self.layer1 = self._make_res_block(filter_num=64, 81 | blocks=block_num[0]) 82 | self.layer2 = self._make_res_block(filter_num=128, 83 | blocks=block_num[1], 84 | stride=2) 85 | self.layer3 = self._make_res_block(filter_num=256, 86 | blocks=block_num[2], 87 | stride=2) 88 | self.layer4 = self._make_res_block(filter_num=512, 89 | blocks=block_num[3], 90 | stride=2) 91 | 92 | self.avgpool = tf.keras.layers.GlobalAveragePooling2D() 93 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, activation=tf.keras.activations.softmax) 94 | 95 | def _make_res_block(self, filter_num, blocks, stride=1): 96 | res_block = tf.keras.Sequential() 97 | res_block.add(BottleNeck(filter_num, stride=stride)) 98 | 99 | for _ in range(1, blocks): 100 | res_block.add(BottleNeck(filter_num, stride=1)) 101 | 102 | return res_block 103 | 104 | def call(self, inputs, training=None, mask=None): 105 | pre1 = self.pre1(inputs) 106 | pre2 = self.pre2(pre1, training=training) 107 | pre3 = self.pre3(pre2) 108 | pre4 = self.pre4(pre3) 109 | l1 = self.layer1(pre4, training=training) 110 | l2 = self.layer2(l1, training=training) 111 | l3 = self.layer3(l2, training=training) 112 | l4 = self.layer4(l3, training=training) 113 | avgpool = self.avgpool(l4) 114 | out = self.fc(avgpool) 115 | return out 116 | 117 | def __repr__(self): 118 | return "SE_ResNet_{}".format(self.model_name) 119 | 120 | 121 | def se_resnet_50(): 122 | return SEResNet(block_num=[3, 4, 6, 3], model_name="50") 123 | 124 | 125 | def se_resnet_101(): 126 | return SEResNet(block_num=[3, 4, 23, 3], model_name="101") 127 | 128 | 129 | def se_resnet_152(): 130 | return SEResNet(block_num=[3, 8, 36, 3], model_name="152") -------------------------------------------------------------------------------- /models/se_resnext.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from configuration import NUM_CLASSES 4 | from models.group_convolution import get_group_conv 5 | 6 | 7 | class SEBlock(tf.keras.layers.Layer): 8 | def __init__(self, input_channels, r=16): 9 | super(SEBlock, self).__init__() 10 | self.pool = tf.keras.layers.GlobalAveragePooling2D() 11 | self.fc1 = tf.keras.layers.Dense(units=input_channels // r) 12 | self.fc2 = tf.keras.layers.Dense(units=input_channels) 13 | 14 | def call(self, inputs, **kwargs): 15 | branch = self.pool(inputs) 16 | branch = self.fc1(branch) 17 | branch = tf.nn.relu(branch) 18 | branch = self.fc2(branch) 19 | branch = tf.nn.sigmoid(branch) 20 | branch = tf.expand_dims(input=branch, axis=1) 21 | branch = tf.expand_dims(input=branch, axis=1) 22 | output = tf.keras.layers.multiply(inputs=[inputs, branch]) 23 | return output 24 | 25 | 26 | class BottleNeck(tf.keras.layers.Layer): 27 | def __init__(self, filters, strides, groups): 28 | super(BottleNeck, self).__init__() 29 | self.conv1 = tf.keras.layers.Conv2D(filters=filters, 30 | kernel_size=(1, 1), 31 | strides=1, 32 | padding="same") 33 | self.bn1 = tf.keras.layers.BatchNormalization() 34 | # self.group_conv = tf.keras.layers.Conv2D(filters=filters, 35 | # kernel_size=(3, 3), 36 | # strides=strides, 37 | # padding="same", 38 | # groups=groups) 39 | self.group_conv = get_group_conv(in_channels=filters, 40 | out_channels=filters, 41 | kernel_size=(3, 3), 42 | strides=strides, 43 | padding="same", 44 | groups=groups) 45 | self.bn2 = tf.keras.layers.BatchNormalization() 46 | self.conv2 = tf.keras.layers.Conv2D(filters=2 * filters, 47 | kernel_size=(1, 1), 48 | strides=1, 49 | padding="same") 50 | self.bn3 = tf.keras.layers.BatchNormalization() 51 | self.se = SEBlock(input_channels=2 * filters) 52 | 53 | self.shortcut_conv = tf.keras.layers.Conv2D(filters=2 * filters, 54 | kernel_size=(1, 1), 55 | strides=strides, 56 | padding="same") 57 | self.shortcut_bn = tf.keras.layers.BatchNormalization() 58 | 59 | def call(self, inputs, training=None, **kwargs): 60 | x = self.conv1(inputs) 61 | x = self.bn1(x, training=training) 62 | x = tf.nn.relu(x) 63 | x = self.group_conv(x) 64 | x = self.bn2(x, training=training) 65 | x = tf.nn.relu(x) 66 | x = self.conv2(x) 67 | x = self.bn3(x, training=training) 68 | x = self.se(x) 69 | 70 | shortcut = self.shortcut_conv(inputs) 71 | shortcut = self.shortcut_bn(shortcut, training=training) 72 | 73 | output = tf.nn.relu(tf.keras.layers.add([x, shortcut])) 74 | return output 75 | 76 | 77 | class SEResNeXt(tf.keras.Model): 78 | def __init__(self, repeat_num_list, cardinality, model_name): 79 | super(SEResNeXt, self).__init__() 80 | self.model_name = model_name 81 | 82 | self.conv1 = tf.keras.layers.Conv2D(filters=64, 83 | kernel_size=(7, 7), 84 | strides=2, 85 | padding="same") 86 | self.bn1 = tf.keras.layers.BatchNormalization() 87 | self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 88 | strides=2, 89 | padding="same") 90 | self.block1 = SEResNeXt.__make_layer(filters=128, 91 | strides=1, 92 | groups=cardinality, 93 | repeat_num=repeat_num_list[0]) 94 | self.block2 = SEResNeXt.__make_layer(filters=256, 95 | strides=2, 96 | groups=cardinality, 97 | repeat_num=repeat_num_list[1]) 98 | self.block3 = SEResNeXt.__make_layer(filters=512, 99 | strides=2, 100 | groups=cardinality, 101 | repeat_num=repeat_num_list[2]) 102 | self.block4 = SEResNeXt.__make_layer(filters=1024, 103 | strides=2, 104 | groups=cardinality, 105 | repeat_num=repeat_num_list[3]) 106 | self.pool2 = tf.keras.layers.GlobalAveragePooling2D() 107 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, 108 | activation=tf.keras.activations.softmax) 109 | 110 | @staticmethod 111 | def __make_layer(filters, strides, groups, repeat_num): 112 | block = tf.keras.Sequential() 113 | block.add(BottleNeck(filters=filters, 114 | strides=strides, 115 | groups=groups)) 116 | for _ in range(1, repeat_num): 117 | block.add(BottleNeck(filters=filters, 118 | strides=1, 119 | groups=groups)) 120 | 121 | return block 122 | 123 | def call(self, inputs, training=None, mask=None): 124 | x = self.conv1(inputs) 125 | x = self.bn1(x, training=training) 126 | x = tf.nn.relu(x) 127 | x = self.pool1(x) 128 | 129 | x = self.block1(x, training=training) 130 | x = self.block2(x, training=training) 131 | x = self.block3(x, training=training) 132 | x = self.block4(x, training=training) 133 | 134 | x = self.pool2(x) 135 | x = self.fc(x) 136 | 137 | return x 138 | 139 | def __repr__(self): 140 | return "SE_ResNeXt_{}".format(self.model_name) 141 | 142 | 143 | def SEResNeXt50(): 144 | return SEResNeXt(repeat_num_list=[3, 4, 6, 3], cardinality=32, model_name="50") 145 | 146 | 147 | def SEResNeXt101(): 148 | return SEResNeXt(repeat_num_list=[3, 4, 23, 3], cardinality=32, model_name="101") -------------------------------------------------------------------------------- /models/shufflenet_v2.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from configuration import NUM_CLASSES 4 | 5 | 6 | def channel_shuffle(feature, group): 7 | channel_num = feature.shape[-1] 8 | if channel_num % group != 0: 9 | raise ValueError("The group must be divisible by the shape of the last dimension of the feature.") 10 | x = tf.reshape(feature, shape=(-1, feature.shape[1], feature.shape[2], group, channel_num // group)) 11 | x = tf.transpose(x, perm=[0, 1, 2, 4, 3]) 12 | x = tf.reshape(x, shape=(-1, feature.shape[1], feature.shape[2], channel_num)) 13 | return x 14 | 15 | 16 | class ShuffleBlockS1(tf.keras.layers.Layer): 17 | def __init__(self, in_channels, out_channels): 18 | super(ShuffleBlockS1, self).__init__() 19 | self.conv1 = tf.keras.layers.Conv2D(filters=out_channels // 2, 20 | kernel_size=(1, 1), 21 | strides=1, 22 | padding="same") 23 | self.bn1 = tf.keras.layers.BatchNormalization() 24 | self.dwconv = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3), strides=1, padding="same") 25 | self.dw_bn = tf.keras.layers.BatchNormalization() 26 | self.conv2 = tf.keras.layers.Conv2D(filters=out_channels // 2, 27 | kernel_size=(1, 1), 28 | strides=1, 29 | padding="same") 30 | self.bn2 = tf.keras.layers.BatchNormalization() 31 | 32 | def call(self, inputs, training=None, **kwargs): 33 | branch, x = tf.split(inputs, num_or_size_splits=2, axis=-1) 34 | x = self.conv1(x) 35 | x = self.bn1(x, training=training) 36 | x = tf.nn.relu(x) 37 | x = self.dwconv(x) 38 | x = self.dw_bn(x, training=training) 39 | x = self.conv2(x) 40 | x = self.bn2(x, training=training) 41 | x = tf.nn.relu(x) 42 | 43 | outputs = tf.concat(values=[branch, x], axis=-1) 44 | outputs = channel_shuffle(feature=outputs, group=2) 45 | return outputs 46 | 47 | 48 | class ShuffleBlockS2(tf.keras.layers.Layer): 49 | def __init__(self, in_channels, out_channels): 50 | super(ShuffleBlockS2, self).__init__() 51 | self.conv1 = tf.keras.layers.Conv2D(filters=out_channels // 2, 52 | kernel_size=(1, 1), 53 | strides=1, 54 | padding="same") 55 | self.bn1 = tf.keras.layers.BatchNormalization() 56 | self.dwconv = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3), strides=2, padding="same") 57 | self.dw_bn = tf.keras.layers.BatchNormalization() 58 | self.conv2 = tf.keras.layers.Conv2D(filters=out_channels - in_channels, 59 | kernel_size=(1, 1), 60 | strides=1, 61 | padding="same") 62 | self.bn2 = tf.keras.layers.BatchNormalization() 63 | 64 | self.branch_dwconv = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3), strides=2, padding="same") 65 | self.branch_dwbn = tf.keras.layers.BatchNormalization() 66 | self.branch_conv = tf.keras.layers.Conv2D(filters=in_channels, 67 | kernel_size=(1, 1), 68 | strides=1, 69 | padding="same") 70 | self.branch_bn = tf.keras.layers.BatchNormalization() 71 | 72 | def call(self, inputs, training=None, **kwargs): 73 | x = self.conv1(inputs) 74 | x = self.bn1(x, training=training) 75 | x = tf.nn.relu(x) 76 | x = self.dwconv(x) 77 | x = self.dw_bn(x, training=training) 78 | x = self.conv2(x) 79 | x = self.bn2(x, training=training) 80 | x = tf.nn.relu(x) 81 | 82 | branch = self.branch_dwconv(inputs) 83 | branch = self.branch_dwbn(branch, training=training) 84 | branch = self.branch_conv(branch) 85 | branch = self.branch_bn(branch, training=training) 86 | branch = tf.nn.relu(branch) 87 | 88 | outputs = tf.concat(values=[x, branch], axis=-1) 89 | outputs = channel_shuffle(feature=outputs, group=2) 90 | return outputs 91 | 92 | 93 | class ShuffleNetV2(tf.keras.Model): 94 | def __init__(self, channel_scale, model_name): 95 | super(ShuffleNetV2, self).__init__() 96 | self.model_name = model_name 97 | 98 | self.conv1 = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3), strides=2, padding="same") 99 | self.bn1 = tf.keras.layers.BatchNormalization() 100 | self.maxpool = tf.keras.layers.MaxPool2D(pool_size=(3, 3), strides=2, padding="same") 101 | self.stage1 = self._make_layer(repeat_num=4, in_channels=24, out_channels=channel_scale[0]) 102 | self.stage2 = self._make_layer(repeat_num=8, in_channels=channel_scale[0], out_channels=channel_scale[1]) 103 | self.stage3 = self._make_layer(repeat_num=4, in_channels=channel_scale[1], out_channels=channel_scale[2]) 104 | self.conv5 = tf.keras.layers.Conv2D(filters=channel_scale[3], kernel_size=(1, 1), strides=1, padding="same") 105 | self.bn5 = tf.keras.layers.BatchNormalization() 106 | self.avgpool = tf.keras.layers.GlobalAveragePooling2D() 107 | self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, activation=tf.keras.activations.softmax) 108 | 109 | def _make_layer(self, repeat_num, in_channels, out_channels): 110 | block = tf.keras.Sequential() 111 | block.add(ShuffleBlockS2(in_channels=in_channels, out_channels=out_channels)) 112 | for i in range(1, repeat_num): 113 | block.add(ShuffleBlockS1(in_channels=out_channels, out_channels=out_channels)) 114 | return block 115 | 116 | def call(self, inputs, training=None, mask=None): 117 | x = self.conv1(inputs) 118 | x = self.bn1(x, training=training) 119 | x = tf.nn.relu(x) 120 | x = self.maxpool(x) 121 | x = self.stage1(x, training=training) 122 | x = self.stage2(x, training=training) 123 | x = self.stage3(x, training=training) 124 | x = self.conv5(x) 125 | x = self.bn5(x, training=training) 126 | x = tf.nn.relu(x) 127 | x = self.avgpool(x) 128 | x = self.fc(x) 129 | return x 130 | 131 | def __repr__(self): 132 | return "ShuffleNetV2_{}".format(self.model_name) 133 | 134 | 135 | def shufflenet_0_5x(): 136 | return ShuffleNetV2(channel_scale=[48, 96, 192, 1024], model_name="0.5x") 137 | 138 | 139 | def shufflenet_1_0x(): 140 | return ShuffleNetV2(channel_scale=[116, 232, 464, 1024], model_name="1.0x") 141 | 142 | 143 | def shufflenet_1_5x(): 144 | return ShuffleNetV2(channel_scale=[176, 352, 704, 1024], model_name="1.5x") 145 | 146 | 147 | def shufflenet_2_0x(): 148 | return ShuffleNetV2(channel_scale=[244, 488, 976, 2048], model_name="2.0x") 149 | -------------------------------------------------------------------------------- /models/squeezenet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from configuration import NUM_CLASSES 3 | 4 | 5 | class FireModule(tf.keras.layers.Layer): 6 | def __init__(self, s1, e1, e3): 7 | super(FireModule, self).__init__() 8 | self.squeeze_layer = tf.keras.layers.Conv2D(filters=s1, 9 | kernel_size=(1, 1), 10 | strides=1, 11 | padding="same") 12 | self.expand_1x1 = tf.keras.layers.Conv2D(filters=e1, 13 | kernel_size=(1, 1), 14 | strides=1, 15 | padding="same") 16 | self.expand_3x3 = tf.keras.layers.Conv2D(filters=e3, 17 | kernel_size=(3, 3), 18 | strides=1, 19 | padding="same") 20 | 21 | def call(self, inputs, **kwargs): 22 | x = self.squeeze_layer(inputs) 23 | x = tf.nn.relu(x) 24 | y1 = self.expand_1x1(x) 25 | y1 = tf.nn.relu(y1) 26 | y2 = self.expand_3x3(x) 27 | y2 = tf.nn.relu(y2) 28 | return tf.concat(values=[y1, y2], axis=-1) 29 | 30 | 31 | class SqueezeNet(tf.keras.Model): 32 | def __init__(self): 33 | super(SqueezeNet, self).__init__() 34 | self.conv1 = tf.keras.layers.Conv2D(filters=96, 35 | kernel_size=(7, 7), 36 | strides=2, 37 | padding="same") 38 | self.maxpool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 39 | strides=2) 40 | self.fire2 = FireModule(s1=16, e1=64, e3=64) 41 | self.fire3 = FireModule(s1=16, e1=64, e3=64) 42 | self.fire4 = FireModule(s1=32, e1=128, e3=128) 43 | self.maxpool4 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 44 | strides=2) 45 | self.fire5 = FireModule(s1=32, e1=128, e3=128) 46 | self.fire6 = FireModule(s1=48, e1=192, e3=192) 47 | self.fire7 = FireModule(s1=48, e1=192, e3=192) 48 | self.fire8 = FireModule(s1=64, e1=256, e3=256) 49 | self.maxpool8 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), 50 | strides=2) 51 | self.fire9 = FireModule(s1=64, e1=256, e3=256) 52 | self.dropout = tf.keras.layers.Dropout(rate=0.5) 53 | self.conv10 = tf.keras.layers.Conv2D(filters=NUM_CLASSES, 54 | kernel_size=(1, 1), 55 | strides=1, 56 | padding="same") 57 | self.avgpool10 = tf.keras.layers.GlobalAveragePooling2D() 58 | 59 | def call(self, inputs, training=None, mask=None): 60 | x = self.conv1(inputs) 61 | x = self.maxpool1(x) 62 | x = self.fire2(x) 63 | x = self.fire3(x) 64 | x = self.fire4(x) 65 | x = self.maxpool4(x) 66 | x = self.fire5(x) 67 | x = self.fire6(x) 68 | x = self.fire7(x) 69 | x = self.fire8(x) 70 | x = self.maxpool8(x) 71 | x = self.fire9(x) 72 | x = self.dropout(x, training=training) 73 | x = self.conv10(x) 74 | x = self.avgpool10(x) 75 | 76 | return tf.nn.softmax(x) 77 | 78 | -------------------------------------------------------------------------------- /original_dataset/README.md: -------------------------------------------------------------------------------- 1 | Please put your pictures for classification here.
2 | The folder name is the type of the pictures which belong to the folder. -------------------------------------------------------------------------------- /parse_tfrecord.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | 4 | def _parse_image_function(example_proto): 5 | # Parse the input tf.Example proto. 6 | return tf.io.parse_single_example(example_proto, { 7 | 'label': tf.io.FixedLenFeature([], tf.dtypes.int64), 8 | 'image_raw': tf.io.FixedLenFeature([], tf.dtypes.string), 9 | }) 10 | 11 | 12 | def get_parsed_dataset(tfrecord_name): 13 | raw_dataset = tf.data.TFRecordDataset(tfrecord_name) 14 | parsed_dataset = raw_dataset.map(_parse_image_function,num_parallel_calls=tf.data.AUTOTUNE) 15 | 16 | return parsed_dataset 17 | -------------------------------------------------------------------------------- /predict.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from configuration import save_model_dir, test_image_dir 3 | from prepare_data import load_and_preprocess_image 4 | from train import get_model 5 | 6 | 7 | def get_single_picture_prediction(model, picture_dir): 8 | image_tensor = load_and_preprocess_image(tf.io.read_file(filename=picture_dir), data_augmentation=False) 9 | image = tf.expand_dims(image_tensor, axis=0) 10 | prediction = model(image, training=False) 11 | pred_class = tf.math.argmax(prediction, axis=-1) 12 | return pred_class 13 | 14 | 15 | if __name__ == '__main__': 16 | # GPU settings 17 | gpus = tf.config.list_physical_devices('GPU') 18 | if gpus: 19 | for gpu in gpus: 20 | tf.config.experimental.set_memory_growth(gpu, True) 21 | 22 | # load the model 23 | model = get_model() 24 | model.load_weights(filepath=save_model_dir+"model") 25 | 26 | pred_class = get_single_picture_prediction(model, test_image_dir) 27 | print(pred_class) -------------------------------------------------------------------------------- /prepare_data.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import pathlib 3 | from configuration import IMAGE_HEIGHT, IMAGE_WIDTH, CHANNELS, \ 4 | BATCH_SIZE, train_tfrecord, valid_tfrecord, test_tfrecord 5 | from parse_tfrecord import get_parsed_dataset 6 | 7 | 8 | def load_and_preprocess_image(image_raw, data_augmentation=False): 9 | # decode 10 | image_tensor = tf.io.decode_image(contents=image_raw, channels=CHANNELS, dtype=tf.dtypes.float32) 11 | 12 | if data_augmentation: 13 | image = tf.image.random_flip_left_right(image=image_tensor) 14 | image = tf.image.resize_with_crop_or_pad(image=image, 15 | target_height=int(IMAGE_HEIGHT * 1.2), 16 | target_width=int(IMAGE_WIDTH * 1.2)) 17 | image = tf.image.random_crop(value=image, size=[IMAGE_HEIGHT, IMAGE_WIDTH, CHANNELS]) 18 | image = tf.image.random_brightness(image=image, max_delta=0.5) 19 | else: 20 | image = tf.image.resize(image_tensor, [IMAGE_HEIGHT, IMAGE_WIDTH]) 21 | 22 | return image 23 | 24 | 25 | def get_images_and_labels(data_root_dir): 26 | # get all images' paths (format: string) 27 | data_root = pathlib.Path(data_root_dir) 28 | all_image_path = [str(path) for path in list(data_root.glob('*/*'))] 29 | # get labels' names 30 | label_names = sorted(item.name for item in data_root.glob('*/')) 31 | # dict: {label : index} 32 | label_to_index = dict((label, index) for index, label in enumerate(label_names)) 33 | # get all images' labels 34 | all_image_label = [label_to_index[pathlib.Path(single_image_path).parent.name] for single_image_path in all_image_path] 35 | 36 | return all_image_path, all_image_label 37 | 38 | 39 | def get_the_length_of_dataset(dataset): 40 | count = 0 41 | for i in dataset: 42 | count += 1 43 | return count 44 | 45 | 46 | def generate_datasets(): 47 | train_dataset = get_parsed_dataset(tfrecord_name=train_tfrecord) 48 | valid_dataset = get_parsed_dataset(tfrecord_name=valid_tfrecord) 49 | test_dataset = get_parsed_dataset(tfrecord_name=test_tfrecord) 50 | 51 | train_count = get_the_length_of_dataset(train_dataset) 52 | valid_count = get_the_length_of_dataset(valid_dataset) 53 | test_count = get_the_length_of_dataset(test_dataset) 54 | 55 | # read the dataset in the form of batch 56 | train_dataset = train_dataset.batch(batch_size=BATCH_SIZE) 57 | valid_dataset = valid_dataset.batch(batch_size=BATCH_SIZE) 58 | test_dataset = test_dataset.batch(batch_size=BATCH_SIZE) 59 | 60 | return train_dataset, valid_dataset, test_dataset, train_count, valid_count, test_count 61 | -------------------------------------------------------------------------------- /saved_model/README.md: -------------------------------------------------------------------------------- 1 | The trained model will be saved here. -------------------------------------------------------------------------------- /show_model_list.py: -------------------------------------------------------------------------------- 1 | from models import get_model2idx_dict 2 | 3 | if __name__ == '__main__': 4 | print("Models included: ") 5 | with open(file="models.txt", mode="w", encoding="utf-8") as f: 6 | for model_name, i in get_model2idx_dict().items(): 7 | print("MODEL: {}, INDEX: {}".format(model_name, i)) 8 | f.write("{}: {}\n".format(model_name, i)) 9 | -------------------------------------------------------------------------------- /split_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import shutil 4 | from configuration import TRAIN_SET_RATIO, TEST_SET_RATIO 5 | 6 | class SplitDataset(): 7 | def __init__(self, dataset_dir, saved_dataset_dir, train_ratio=TRAIN_SET_RATIO, test_ratio=TEST_SET_RATIO, show_progress=False): 8 | self.dataset_dir = dataset_dir 9 | self.saved_dataset_dir = saved_dataset_dir 10 | self.saved_train_dir = saved_dataset_dir + "/train/" 11 | self.saved_valid_dir = saved_dataset_dir + "/valid/" 12 | self.saved_test_dir = saved_dataset_dir + "/test/" 13 | 14 | 15 | self.train_ratio = train_ratio 16 | self.test_radio = test_ratio 17 | self.valid_ratio = 1 - train_ratio - test_ratio 18 | 19 | self.train_file_path = [] 20 | self.valid_file_path = [] 21 | self.test_file_path = [] 22 | 23 | self.index_label_dict = {} 24 | 25 | self.show_progress = show_progress 26 | 27 | if not os.path.exists(self.saved_train_dir): 28 | os.mkdir(self.saved_train_dir) 29 | if not os.path.exists(self.saved_test_dir): 30 | os.mkdir(self.saved_test_dir) 31 | if not os.path.exists(self.saved_valid_dir): 32 | os.mkdir(self.saved_valid_dir) 33 | 34 | 35 | def __get_label_names(self): 36 | label_names = [] 37 | for item in os.listdir(self.dataset_dir): 38 | item_path = os.path.join(self.dataset_dir, item) 39 | if os.path.isdir(item_path): 40 | label_names.append(item) 41 | return label_names 42 | 43 | def __get_all_file_path(self): 44 | all_file_path = [] 45 | index = 0 46 | for file_type in self.__get_label_names(): 47 | self.index_label_dict[index] = file_type 48 | index += 1 49 | type_file_path = os.path.join(self.dataset_dir, file_type) 50 | file_path = [] 51 | for file in os.listdir(type_file_path): 52 | single_file_path = os.path.join(type_file_path, file) 53 | file_path.append(single_file_path) 54 | all_file_path.append(file_path) 55 | return all_file_path 56 | 57 | def __copy_files(self, type_path, type_saved_dir): 58 | for item in type_path: 59 | src_path_list = item[1] 60 | dst_path = type_saved_dir + "%s/" % (item[0]) 61 | if not os.path.exists(dst_path): 62 | os.mkdir(dst_path) 63 | for src_path in src_path_list: 64 | shutil.copy(src_path, dst_path) 65 | if self.show_progress: 66 | print("Copying file "+src_path+" to "+dst_path) 67 | 68 | def __split_dataset(self): 69 | all_file_paths = self.__get_all_file_path() 70 | for index in range(len(all_file_paths)): 71 | file_path_list = all_file_paths[index] 72 | file_path_list_length = len(file_path_list) 73 | random.shuffle(file_path_list) 74 | 75 | train_num = int(file_path_list_length * self.train_ratio) 76 | test_num = int(file_path_list_length * self.test_radio) 77 | 78 | self.train_file_path.append([self.index_label_dict[index], file_path_list[: train_num]]) 79 | self.test_file_path.append([self.index_label_dict[index], file_path_list[train_num:train_num + test_num]]) 80 | self.valid_file_path.append([self.index_label_dict[index], file_path_list[train_num + test_num:]]) 81 | 82 | def start_splitting(self): 83 | self.__split_dataset() 84 | self.__copy_files(type_path=self.train_file_path, type_saved_dir=self.saved_train_dir) 85 | self.__copy_files(type_path=self.valid_file_path, type_saved_dir=self.saved_valid_dir) 86 | self.__copy_files(type_path=self.test_file_path, type_saved_dir=self.saved_test_dir) 87 | 88 | 89 | if __name__ == '__main__': 90 | split_dataset = SplitDataset(dataset_dir="original_dataset", 91 | saved_dataset_dir="dataset", 92 | show_progress=True) 93 | split_dataset.start_splitting() -------------------------------------------------------------------------------- /test_single_image.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import os 3 | 4 | from configuration import save_model_dir, test_image_dir 5 | from train import get_model 6 | from prepare_data import load_and_preprocess_image 7 | 8 | 9 | def get_class_id(image_root): 10 | id_cls = {} 11 | for i, item in enumerate(os.listdir(image_root)): 12 | if os.path.isdir(os.path.join(image_root, item)): 13 | id_cls[i] = item 14 | return id_cls 15 | 16 | 17 | if __name__ == '__main__': 18 | # GPU settings 19 | gpus = tf.config.list_physical_devices('GPU') 20 | if gpus: 21 | for gpu in gpus: 22 | tf.config.experimental.set_memory_growth(gpu, True) 23 | 24 | model = get_model() 25 | model.load_weights(filepath=save_model_dir) 26 | 27 | image_raw = tf.io.read_file(filename=test_image_dir) 28 | image_tensor = load_and_preprocess_image(image_raw) 29 | image_tensor = tf.expand_dims(image_tensor, axis=0) 30 | 31 | pred = model(image_tensor, training=False) 32 | idx = tf.math.argmax(pred, axis=-1).numpy()[0] 33 | 34 | id_cls = get_class_id("./original_dataset") 35 | 36 | print("The predicted category of this picture is: {}".format(id_cls[idx])) -------------------------------------------------------------------------------- /to_tfrecord.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from configuration import train_dir, valid_dir, test_dir, train_tfrecord, valid_tfrecord, test_tfrecord 3 | from prepare_data import get_images_and_labels 4 | import random 5 | 6 | # convert a value to a type compatible tf.train.Feature 7 | def _bytes_feature(value): 8 | # Returns a bytes_list from a string / byte. 9 | if isinstance(value, type(tf.constant(0.))): 10 | value = value.numpy() # BytesList won't unpack a string from an EagerTensor. 11 | return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 12 | 13 | 14 | def _float_feature(value): 15 | # Returns a float_list from a float / double. 16 | return tf.train.Feature(float_list=tf.train.FloatList(value=[value])) 17 | 18 | 19 | def _int64_feature(value): 20 | # Returns an int64_list from a bool / enum / int / uint. 21 | return tf.train.Feature(int64_list=tf.train.Int64List(value=[value])) 22 | 23 | 24 | # Create a dictionary with features that may be relevant. 25 | def image_example(image_string, label): 26 | feature = { 27 | 'label': _int64_feature(label), 28 | 'image_raw': _bytes_feature(image_string), 29 | } 30 | 31 | return tf.train.Example(features=tf.train.Features(feature=feature)) 32 | 33 | 34 | def shuffle_dict(original_dict): 35 | keys = [] 36 | shuffled_dict = {} 37 | for k in original_dict.keys(): 38 | keys.append(k) 39 | random.shuffle(keys) 40 | for item in keys: 41 | shuffled_dict[item] = original_dict[item] 42 | return shuffled_dict 43 | 44 | 45 | def dataset_to_tfrecord(dataset_dir, tfrecord_name): 46 | image_paths, image_labels = get_images_and_labels(dataset_dir) 47 | image_paths_and_labels_dict = {} 48 | for i in range(len(image_paths)): 49 | image_paths_and_labels_dict[image_paths[i]] = image_labels[i] 50 | # shuffle the dict 51 | image_paths_and_labels_dict = shuffle_dict(image_paths_and_labels_dict) 52 | # write the images and labels to tfrecord format file 53 | with tf.io.TFRecordWriter(path=tfrecord_name) as writer: 54 | for image_path, label in image_paths_and_labels_dict.items(): 55 | print("Writing to tfrecord: {}".format(image_path)) 56 | image_string = open(image_path, 'rb').read() 57 | tf_example = image_example(image_string, label) 58 | writer.write(tf_example.SerializeToString()) 59 | 60 | 61 | if __name__ == '__main__': 62 | dataset_to_tfrecord(dataset_dir=train_dir, tfrecord_name=train_tfrecord) 63 | dataset_to_tfrecord(dataset_dir=valid_dir, tfrecord_name=valid_tfrecord) 64 | dataset_to_tfrecord(dataset_dir=test_dir, tfrecord_name=test_tfrecord) -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function 2 | import tensorflow as tf 3 | import tensorflow.keras as nn 4 | import math 5 | import argparse 6 | 7 | from configuration import IMAGE_HEIGHT, IMAGE_WIDTH, CHANNELS, \ 8 | EPOCHS, BATCH_SIZE, save_model_dir, save_every_n_epoch 9 | from prepare_data import generate_datasets, load_and_preprocess_image 10 | from models import get_model 11 | 12 | 13 | def print_model_summary(network): 14 | network.build(input_shape=(None, IMAGE_HEIGHT, IMAGE_WIDTH, CHANNELS)) 15 | network.summary() 16 | 17 | 18 | def process_features(features, data_augmentation): 19 | image_raw = features['image_raw'].numpy() 20 | image_tensor_list = [] 21 | for image in image_raw: 22 | image_tensor = load_and_preprocess_image(image, data_augmentation=data_augmentation) 23 | image_tensor_list.append(image_tensor) 24 | images = tf.stack(image_tensor_list, axis=0) 25 | labels = features['label'].numpy() 26 | 27 | return images, labels 28 | 29 | 30 | parser = argparse.ArgumentParser() 31 | parser.add_argument("--idx", default=0, type=int) 32 | 33 | 34 | if __name__ == '__main__': 35 | gpus = tf.config.list_physical_devices('GPU') 36 | if gpus: 37 | try: 38 | for gpu in gpus: 39 | tf.config.experimental.set_memory_growth(gpu, True) 40 | logical_gpus = tf.config.list_logical_devices('GPU') 41 | print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") 42 | except RuntimeError as e: 43 | print(e) 44 | 45 | args = parser.parse_args() 46 | 47 | # get the dataset 48 | train_dataset, valid_dataset, test_dataset, train_count, valid_count, test_count = generate_datasets() 49 | 50 | # create model 51 | model = get_model(args.idx) 52 | print_model_summary(network=model) 53 | 54 | # define loss and optimizer 55 | loss_object = tf.keras.losses.SparseCategoricalCrossentropy() 56 | optimizer = nn.optimizers.Adam(learning_rate=1e-3) 57 | 58 | train_loss = tf.keras.metrics.Mean(name='train_loss') 59 | train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy') 60 | 61 | valid_loss = tf.keras.metrics.Mean(name='valid_loss') 62 | valid_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='valid_accuracy') 63 | 64 | # @tf.function 65 | def train_step(image_batch, label_batch): 66 | with tf.GradientTape() as tape: 67 | predictions = model(image_batch, training=True) 68 | loss = loss_object(y_true=label_batch, y_pred=predictions) 69 | gradients = tape.gradient(loss, model.trainable_variables) 70 | optimizer.apply_gradients(grads_and_vars=zip(gradients, model.trainable_variables)) 71 | 72 | train_loss.update_state(values=loss) 73 | train_accuracy.update_state(y_true=label_batch, y_pred=predictions) 74 | 75 | # @tf.function 76 | def valid_step(image_batch, label_batch): 77 | predictions = model(image_batch, training=False) 78 | v_loss = loss_object(label_batch, predictions) 79 | 80 | valid_loss.update_state(values=v_loss) 81 | valid_accuracy.update_state(y_true=label_batch, y_pred=predictions) 82 | 83 | # start training 84 | for epoch in range(EPOCHS): 85 | step = 0 86 | for features in train_dataset: 87 | step += 1 88 | images, labels = process_features(features, data_augmentation=True) 89 | train_step(images, labels) 90 | print("Epoch: {}/{}, step: {}/{}, loss: {:.5f}, accuracy: {:.5f}".format(epoch, 91 | EPOCHS, 92 | step, 93 | math.ceil(train_count / BATCH_SIZE), 94 | train_loss.result().numpy(), 95 | train_accuracy.result().numpy())) 96 | 97 | for features in valid_dataset: 98 | valid_images, valid_labels = process_features(features, data_augmentation=False) 99 | valid_step(valid_images, valid_labels) 100 | 101 | print("Epoch: {}/{}, train loss: {:.5f}, train accuracy: {:.5f}, " 102 | "valid loss: {:.5f}, valid accuracy: {:.5f}".format(epoch, 103 | EPOCHS, 104 | train_loss.result().numpy(), 105 | train_accuracy.result().numpy(), 106 | valid_loss.result().numpy(), 107 | valid_accuracy.result().numpy())) 108 | train_loss.reset_states() 109 | train_accuracy.reset_states() 110 | valid_loss.reset_states() 111 | valid_accuracy.reset_states() 112 | 113 | if epoch % save_every_n_epoch == 0: 114 | model.save_weights(filepath=save_model_dir+"epoch-{}".format(epoch), save_format='tf') 115 | 116 | 117 | # save weights 118 | model.save_weights(filepath=save_model_dir+"model", save_format='tf') 119 | 120 | # save the whole model 121 | # tf.saved_model.save(model, save_model_dir) 122 | 123 | # convert to tensorflow lite format 124 | # model._set_inputs(inputs=tf.random.normal(shape=(1, IMAGE_HEIGHT, IMAGE_WIDTH, CHANNELS))) 125 | # converter = tf.lite.TFLiteConverter.from_keras_model(model) 126 | # tflite_model = converter.convert() 127 | # open("converted_model.tflite", "wb").write(tflite_model) 128 | 129 | --------------------------------------------------------------------------------