classifier-fine-tuning.py

#!/usr/bin/env python3
import tensorflow as tf
import pathlib


data_dir = pathlib.Path.home() / 'datasets' / 'dogs-v-cats' / 'train'
# data_dir = pathlib.Path.home() / 'datasets' / '10flowers' / 'images'
batch_size = 32
img_height = 224
img_width = 224
epochs = 10
use_pretrained_model = False


print('Images:', len(list(data_dir.glob('*.jpg'))))
train_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    labels='inferred',
    validation_split=0.2,
    subset='training',
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size
)
val_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    labels='inferred',
    validation_split=0.2,
    subset='validation',
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size
)
class_names = train_ds.class_names


@tf.function
def cast_x_to_float(x, y):
    return tf.cast(x, tf.float32), y


train_ds = train_ds.map(cast_x_to_float)
val_ds = val_ds.map(cast_x_to_float)


# `cache()` means the dataset gets stored in memory as it's read so it
# doesn't have to be read multiple times from the hard drive (which is
# slower than the main system memory)
# in an ideal world, the dataset could all be stored in the GPU, but usually
# it's a better idea to reserve that precious GPU memory for other things
# remember, most time a computer spends doing things is not spent on
# calculations but instead on moving data from HDD to mem to CPU to GPU, etc.
train_ds = train_ds.cache()
val_ds = val_ds.cache()


# ATTENTION YUIGO!
# Let's do two things to improve this machine learning pipeline:
# 1. let's replace the three conv layers with a model that's been
# pretrained on imagenet
# 2. let's do some data augmentation to help avoid overfitting

# here's (2) our data augmentation code, notice I only augment the training set
flip = tf.keras.layers.RandomFlip('horizontal')
rotate = tf.keras.layers.RandomRotation(0.2)


@tf.function
def augment(image, label):
    p = 0.5
    if tf.random.uniform([]) < p:
        image = flip(image)
    if tf.random.uniform([]) < p:
        image = rotate(image)
    return image, label


train_ds = train_ds.map(augment)  # comment-out this line to try without augmentation

# always remember to normalize data so it's values aren't too large in magnitude
# there are many slight variations to this -- since images are stored as uint8
# values (integers 0-255), usually this just means dividing by 255.
# When using a pretrained model it can help a bit to know how a model was normalized
# when it was originally trained.  Sometimes they come with "preprocess" function
# that takes care of this step.  We'll use their function, but it's
# probably very similar to the `a_typical_way_to_normalize` function below.
# Also notice, since we're doing this here, I removed the `tf.keras.layers.Rescaling(1. / 255)`


@tf.function
def normalize(x, y):
    x = tf.keras.applications.mobilenet_v2.preprocess_input(x)
    return x, y


@tf.function
def a_typical_way_to_normalize(x, y):
    """scale images to have values between -0.5 and 0.5 (instead of 0 to 255)"""
    x = tf.cast(x, tf.float32) / 255. - 0.5
    return x, y


train_ds = train_ds.map(normalize)
val_ds = val_ds.map(normalize)


# here's (1) our new model
# there are many pre-trained models you can easily use with the keras api
# see https://www.tensorflow.org/api_docs/python/tf/keras/applications
if use_pretrained_model:
    model = tf.keras.Sequential([
        tf.keras.applications.MobileNetV2(
            input_shape=(img_height, img_width, 3),
            include_top=False,  # this means the dense layers at the end are removed
            weights='imagenet'
        ),
        # here we add back (untrained) dense layers to replace the ones removed
        tf.keras.layers.Flatten(),
        # tf.keras.layers.Dense(64, activation='relu'),  # this extra layer may help or may hurt
        tf.keras.layers.Dense(len(class_names))
    ])
else:
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(64, activation='relu'),  # this extra layer may help or may hurt
        tf.keras.layers.Dense(len(class_names)),
    ])


model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

# `prefetch()` loads some images to the gpu to before they're needed,
# which can speed things up, or slow them down -- autotune experiments a
# bit and tries to decide the best setting for your system
AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.prefetch(buffer_size=AUTOTUNE)

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