"""Image classification with EfficientNetV2 architecture. Adapted from the EfficientNetV2 Keras Application. """ import math from tensorflow import keras from keras.integration_test.models.input_spec import InputSpec IMG_SIZE = (96, 96) NUM_CLASSES = 5 def get_data_spec(batch_size): return ( InputSpec((batch_size,) + IMG_SIZE + (3,)), InputSpec((batch_size, NUM_CLASSES)), ) def get_input_preprocessor(): return keras.layers.Rescaling(scale=1.0 / 128.0, offset=-1) def round_filters(filters, width_coefficient, min_depth, depth_divisor): filters *= width_coefficient minimum_depth = min_depth or depth_divisor new_filters = max( minimum_depth, int(filters + depth_divisor / 2) // depth_divisor * depth_divisor, ) return int(new_filters) def MBConvBlock( input_filters: int, output_filters: int, expand_ratio=1, kernel_size=3, strides=1, se_ratio=0.0, activation="swish", survival_probability: float = 0.8, ): def apply(inputs): filters = input_filters * expand_ratio if expand_ratio != 1: x = keras.layers.Conv2D( filters=filters, kernel_size=1, strides=1, padding="same", data_format="channels_last", use_bias=False, )(inputs) x = keras.layers.BatchNormalization()(x) x = keras.layers.Activation(activation)(x) else: x = inputs x = keras.layers.DepthwiseConv2D( kernel_size=kernel_size, strides=strides, padding="same", data_format="channels_last", use_bias=False, )(x) x = keras.layers.BatchNormalization()(x) x = keras.layers.Activation(activation)(x) if 0 < se_ratio <= 1: filters_se = max(1, int(input_filters * se_ratio)) se = keras.layers.GlobalAveragePooling2D()(x) se = keras.layers.Reshape((1, 1, filters))(se) se = keras.layers.Conv2D( filters_se, 1, padding="same", activation=activation, )(se) se = keras.layers.Conv2D( filters, 1, padding="same", activation="sigmoid", )(se) x = keras.layers.multiply([x, se]) x = keras.layers.Conv2D( filters=output_filters, kernel_size=1, strides=1, padding="same", data_format="channels_last", use_bias=False, )(x) x = keras.layers.BatchNormalization()(x) if strides == 1 and input_filters == output_filters: if survival_probability: x = keras.layers.Dropout( survival_probability, noise_shape=(None, 1, 1, 1), )(x) x = keras.layers.add([x, inputs]) return x return apply def FusedMBConvBlock( input_filters: int, output_filters: int, expand_ratio=1, kernel_size=3, strides=1, se_ratio=0.0, activation="swish", survival_probability: float = 0.8, ): def apply(inputs): filters = input_filters * expand_ratio if expand_ratio != 1: x = keras.layers.Conv2D( filters, kernel_size=kernel_size, strides=strides, data_format="channels_last", padding="same", use_bias=False, )(inputs) x = keras.layers.BatchNormalization()(x) x = keras.layers.Activation(activation)(x) else: x = inputs if 0 < se_ratio <= 1: filters_se = max(1, int(input_filters * se_ratio)) se = keras.layers.GlobalAveragePooling2D()(x) se = keras.layers.Reshape((1, 1, filters))(se) se = keras.layers.Conv2D( filters_se, 1, padding="same", activation=activation, )(se) se = keras.layers.Conv2D( filters, 1, padding="same", activation="sigmoid", )(se) x = keras.layers.multiply([x, se]) x = keras.layers.Conv2D( output_filters, kernel_size=1 if expand_ratio != 1 else kernel_size, strides=1 if expand_ratio != 1 else strides, padding="same", use_bias=False, )(x) x = keras.layers.BatchNormalization()(x) if expand_ratio == 1: x = keras.layers.Activation(activation)(x) if strides == 1 and input_filters == output_filters: if survival_probability: x = keras.layers.Dropout( survival_probability, noise_shape=(None, 1, 1, 1), )(x) x = keras.layers.add([x, inputs]) return x return apply def get_model( build=False, compile=False, jit_compile=False, include_preprocessing=True ): width_coefficient = 1.0 depth_coefficient = 1.0 dropout_rate = 0.2 drop_connect_rate = 0.2 depth_divisor = 8 min_depth = 8 activation = "swish" blocks_args = [ { "kernel_size": 3, "num_repeat": 2, "input_filters": 24, "output_filters": 24, "expand_ratio": 1, "se_ratio": 0.0, "strides": 1, "conv_type": 1, }, { "kernel_size": 3, "num_repeat": 4, "input_filters": 24, "output_filters": 48, "expand_ratio": 4, "se_ratio": 0.0, "strides": 2, "conv_type": 1, }, { "conv_type": 1, "expand_ratio": 4, "input_filters": 48, "kernel_size": 3, "num_repeat": 4, "output_filters": 64, "se_ratio": 0, "strides": 2, }, { "conv_type": 0, "expand_ratio": 4, "input_filters": 64, "kernel_size": 3, "num_repeat": 6, "output_filters": 128, "se_ratio": 0.25, "strides": 2, }, ] inputs = keras.layers.Input(shape=IMG_SIZE + (3,)) if include_preprocessing: x = get_input_preprocessor()(inputs) else: x = inputs stem_filters = round_filters( filters=blocks_args[0]["input_filters"], width_coefficient=width_coefficient, min_depth=min_depth, depth_divisor=depth_divisor, ) x = keras.layers.Conv2D( filters=stem_filters, kernel_size=3, strides=2, padding="same", use_bias=False, )(x) x = keras.layers.BatchNormalization()(x) x = keras.layers.Activation(activation, name="stem_activation")(x) b = 0 blocks = float(sum(args["num_repeat"] for args in blocks_args)) for _, args in enumerate(blocks_args): args["input_filters"] = round_filters( filters=args["input_filters"], width_coefficient=width_coefficient, min_depth=min_depth, depth_divisor=depth_divisor, ) args["output_filters"] = round_filters( filters=args["output_filters"], width_coefficient=width_coefficient, min_depth=min_depth, depth_divisor=depth_divisor, ) block = {0: MBConvBlock, 1: FusedMBConvBlock}[args.pop("conv_type")] repeats = int(math.ceil(depth_coefficient * args.pop("num_repeat"))) for j in range(repeats): if j > 0: args["strides"] = 1 args["input_filters"] = args["output_filters"] x = block( activation=activation, survival_probability=drop_connect_rate * b / blocks, **args, )(x) b += 1 top_filters = round_filters( filters=1280, width_coefficient=width_coefficient, min_depth=min_depth, depth_divisor=depth_divisor, ) x = keras.layers.Conv2D( filters=top_filters, kernel_size=1, strides=1, padding="same", data_format="channels_last", use_bias=False, )(x) x = keras.layers.BatchNormalization()(x) x = keras.layers.Activation(activation=activation, name="top_activation")(x) x = keras.layers.GlobalAveragePooling2D(name="avg_pool")(x) x = keras.layers.Dropout(dropout_rate, name="top_dropout")(x) x = keras.layers.Dense( NUM_CLASSES, activation="softmax", )(x) model = keras.Model(inputs, x) if compile: model.compile( "adam", loss="categorical_crossentropy", jit_compile=jit_compile ) return model def get_custom_objects(): return {}