Training Image (Binary) Classification with Keras, EfficientNet

efficientnet.py

# import the necessary packages
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.keras.applications.efficientnet import EfficientNetB3, EfficientNetB4, EfficientNetB5, preprocess_input

'''
    EfficientNetB0 - (224, 224, 3)
    EfficientNetB1 - (240, 240, 3)
    EfficientNetB2 - (260, 260, 3)
    EfficientNetB3 - (300, 300, 3)
    EfficientNetB4 - (380, 380, 3)
    EfficientNetB5 - (456, 456, 3)
    EfficientNetB6 - (528, 528, 3)
    EfficientNetB7 - (600, 600, 3)
''' 
class OwnEfficientNetB3:
    @staticmethod
    def build(input_shape, data_augmentation, trainable=False, dropout=0.2):
        inputs = keras.Input(shape=input_shape)
        x = data_augmentation(inputs)
        x = preprocess_input(x)

        baseModel = EfficientNetB3(weights="imagenet", include_top=False, input_tensor=x)
        baseModel.trainable = trainable

        headModel = baseModel.output
        headModel = layers.GlobalAveragePooling2D()(headModel)
        headModel = layers.Dropout(dropout)(headModel)
        outputs = layers.Dense(1, activation="sigmoid")(headModel)
        model = Model(inputs, outputs)

        return model

class OwnEfficientNetB4:
    @staticmethod
    def build(input_shape, data_augmentation, trainable=False, dropout=0.2):
        inputs = keras.Input(shape=input_shape)
        x = data_augmentation(inputs)
        x = preprocess_input(x)

        baseModel = EfficientNetB4(weights="imagenet", include_top=False, input_tensor=x)
        baseModel.trainable = trainable

        headModel = baseModel.output
        headModel = layers.GlobalAveragePooling2D()(headModel)
        headModel = layers.Dropout(dropout)(headModel)
        outputs = layers.Dense(1, activation="sigmoid")(headModel)
        model = Model(inputs, outputs)

        return model

class OwnEfficientNetB5:
    @staticmethod
    def build(input_shape, data_augmentation, trainable=False, dropout=0.2):
        inputs = keras.Input(shape=input_shape)
        x = data_augmentation(inputs)
        x = preprocess_input(x)

        baseModel = EfficientNetB5(weights="imagenet", include_top=False, input_tensor=x)
        baseModel.trainable = trainable

        headModel = baseModel.output
        headModel = layers.GlobalAveragePooling2D()(headModel)
        headModel = layers.Dropout(dropout)(headModel)
        outputs = layers.Dense(1, activation="sigmoid")(headModel)
        model = Model(inputs, outputs)

        return model

fine-tune.py

import tensorflow as tf
from tensorflow import keras
from keras import layers
from keras.callbacks import TensorBoard, EarlyStopping, ModelCheckpoint
from utils.keras_hist_graph import plot_history

from time import time
print(tf.__version__)

# Model configuration
data_path = 'data/dresses'
img_width, img_height = 300, 300
image_size = (img_width, img_height)
batch_size = 32
epochs = 500
validation_split = 0.2
verbosity = 1
seed = 168369

model_path = 'checkpoints/effnetb5-save_binary_169-0.10.hdf5'

print("[INFO] loading network...")
model = tf.keras.models.load_model(model_path)

#fine tune
trainable = True
learning_rate = 1e-5
initial_epoch = 169

if initial_epoch >= epochs:
    print(f'Warning: initial_epoch {initial_epoch} is >= total epochs {epochs}')

train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_path,
    validation_split=validation_split,
    subset="training",
    seed=seed,
    image_size=image_size,
    batch_size=batch_size,
)

print(train_ds)

class_names = train_ds.class_names
print(class_names)

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_path,
    validation_split=validation_split,
    subset="validation",
    seed=seed,
    image_size=image_size,
    batch_size=batch_size,
)

#save some io
AUTOTUNE = tf.data.AUTOTUNE

train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

callbacks = [
    EarlyStopping(monitor='val_loss', patience=15, mode='min', min_delta=0.0001),
    ModelCheckpoint("checkpoints/dresses-finetuned-only-batch-layers-save_binary_{epoch:02d}-{val_loss:.2f}.hdf5", monitor='val_loss', mode='min')
]

#keep batchlayers intact, reduces learning time
for layer in model.layers:
    if isinstance(layer, layers.BatchNormalization):
        layer.trainable = False
    else:
        layer.trainable = True

model.compile(optimizer=keras.optimizers.Adam(learning_rate),
            loss=tf.keras.losses.BinaryCrossentropy(),
            metrics=keras.metrics.BinaryAccuracy())
model.summary()

history = model.fit(
    train_ds, epochs=epochs, callbacks=callbacks, validation_data=val_ds, initial_epoch=initial_epoch
)

plot_history(history)

train.py

import tensorflow as tf
from tensorflow import keras
from keras import layers
from keras.callbacks import TensorBoard, EarlyStopping, ModelCheckpoint
from model.efficientnet import *
import argparse

from time import time
print(tf.__version__)

ap = argparse.ArgumentParser()
ap.add_argument("-d", "--data", required=True, help="path to the data used for classification")
ap.add_argument("-w", "--width", default=456, type=int, help="width for the image")
ap.add_argument("-ht", "--height", default=456, type=int, help="height for the image")
ap.add_argument("-b", "--batch_size", default=32, type=int, help="batch size")
ap.add_argument("-e", "--epochs", default=500, type=int, help="number of epochs")
ap.add_argument("-v", "--validation_split", default=0.2, type=float, help="validation split")
ap.add_argument("-s", "--seed", default=168369, type=int, help="seed")

args = vars(ap.parse_args())

data_path = args["data"]
img_width, img_height = args["width"], args["height"]
image_size = (img_width, img_height)
batch_size = args["batch_size"]
epochs = args["epochs"]
validation_split = args["validation_split"]
seed = args["seed"]
verbosity = 1

#fine tune
learning_rate = 1e-3

train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_path,
    validation_split=validation_split,
    subset="training",
    seed=seed,
    image_size=image_size,
    batch_size=batch_size,
)

print(train_ds)

class_names = train_ds.class_names
print(class_names)

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_path,
    validation_split=validation_split,
    subset="validation",
    seed=seed,
    image_size=image_size,
    batch_size=batch_size,
)

data_augmentation = keras.Sequential(
    [
        layers.RandomFlip("horizontal"),
        layers.RandomRotation(0.1),
        layers.RandomZoom(0.1),
        layers.RandomContrast(factor=0.1),
        layers.RandomTranslation(height_factor=0.1, width_factor=0.1),
    ],
    name="img_augmentation"
)

#save some io
AUTOTUNE = tf.data.AUTOTUNE

train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

model = OwnEfficientNetB5.build(input_shape=image_size + (3,), data_augmentation=data_augmentation)

callbacks = [
    EarlyStopping(monitor='val_loss', patience=30, mode='min', min_delta=0.0001),
    ModelCheckpoint("checkpoints/effnetb5-save_binary_{epoch:02d}-{val_loss:.2f}.hdf5", monitor='val_loss', mode='min', save_best_only=True)
]

model.compile(optimizer=keras.optimizers.Adam(learning_rate),
            loss=tf.keras.losses.BinaryCrossentropy(),
            metrics=keras.metrics.BinaryAccuracy())

model.summary()

model.fit(
    train_ds, epochs=epochs, callbacks=callbacks, validation_data=val_ds,
)