Face detect CNN training example

facedetect_train.py

"""
Example adapted from https://github.com/aamini/introtodeeplearning lab2 part2

© MIT Introduction to Deep Learning
http://introtodeeplearning.com
"""
import tensorflow as tf
import matplotlib.pyplot as plt
import functools
import numpy as np
import h5py

# load training data
path_to_training_data = tf.keras.utils.get_file(
    'train_face.h5', 'https://www.dropbox.com/s/hlz8atheyozp1yx/train_face.h5?dl=1')
cache = h5py.File(path_to_training_data, 'r')
images = cache['images'][:]
labels = cache['labels'][:].astype(np.float32)
n_train_samples = images.shape[0]
np.random.seed(4)
train_inds = np.random.permutation(np.arange(n_train_samples))
pos_train_inds = train_inds[labels[train_inds, 0] == 1.0]
neg_train_inds = train_inds[labels[train_inds, 0] != 1.0]


def get_batch(n):
    selected_pos_inds = np.random.choice(
        pos_train_inds, size=n//2, replace=False, p=None)
    selected_neg_inds = np.random.choice(
        neg_train_inds, size=n//2, replace=False, p=None)
    selected_inds = np.concatenate((selected_pos_inds, selected_neg_inds))

    sorted_inds = np.sort(selected_inds)
    train_img = (images[sorted_inds, :, :, ::-1]/255.).astype(np.float32)
    train_label = labels[sorted_inds, ...]
    return train_img, train_label


# define model
n_filters = 12


def make_standard_classifier(n_outputs=1):
    Conv2D = functools.partial(
        tf.keras.layers.Conv2D, padding='same', activation='relu')
    BatchNormalization = tf.keras.layers.BatchNormalization
    Flatten = tf.keras.layers.Flatten
    Dense = functools.partial(tf.keras.layers.Dense, activation='relu')

    model = tf.keras.Sequential([
        Conv2D(filters=1*n_filters, kernel_size=5,  strides=2),
        BatchNormalization(),

        Conv2D(filters=2*n_filters, kernel_size=5,  strides=2),
        BatchNormalization(),

        Conv2D(filters=4*n_filters, kernel_size=3,  strides=2),
        BatchNormalization(),

        Conv2D(filters=6*n_filters, kernel_size=3,  strides=2),
        BatchNormalization(),

        Flatten(),
        Dense(512),
        Dense(n_outputs, activation=None),
    ])
    return model


model = make_standard_classifier()

batch_size = 32
num_epochs = 2
learning_rate = 5e-4
optimizer = tf.keras.optimizers.Adam(learning_rate)
loss_history = []
smooth = 0.99


@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x)
        loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=y, logits=logits)

    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss


for epoch in range(num_epochs):
    print(f"Epoch {epoch}/{num_epochs}")
    for idx in range(n_train_samples//batch_size):
        x, y = get_batch(batch_size)
        loss = train_step(x, y)

        loss_scalar = loss.numpy().mean()
        loss_history.append(
            smooth*loss_history[-1] + (1-smooth)*loss_scalar if len(loss_history) > 0 else loss_scalar)
        if idx % 400 == 0:
            print(f"  batch {idx:4d} loss={loss_history[-1]:.5f}")

plt.semilogy(loss_history)
imgname = 'loss_history.png'
print(f'Saving image {imgname}')
plt.savefig(imgname, dpi=200, bbox_inches='tight')

(batch_x, batch_y) = get_batch(5000)
y_pred = tf.round(tf.nn.sigmoid(model.predict(batch_x)))
acc = tf.reduce_mean(tf.cast(tf.equal(batch_y, y_pred), tf.float32))
print("CNN accuracy on training set: {:.4f}".format(acc.numpy()))