CNN Training Script using TensorFlow using pre-trained VGG-16 weights.

train_tf.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf
import os
import dataset
import math
import random
import pdb

from numpy.random import seed
seed(1)
from tensorflow import set_random_seed
set_random_seed(2)

os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
tf.logging.set_verbosity(tf.logging.INFO)

# VGG16 Weights in npz format
vgg16_weights = np.load(os.path.join(os.path.expanduser('~'),'Downloads/vgg16_weights.npz')) 

batch_size = 8

#Prepare input data
classes = ['c0','c1']
num_classes = len(classes)

fc_layer_size = 256


# 20% of the data will automatically be used for validation
validation_size = 0.2
img_size = 150
num_channels = 3
train_path = os.path.join(os.path.expanduser('~'),'200_image/dataset/') # change the path here

# We shall load all the training and validation images and labels into memory using openCV and use that during training
data = dataset.read_train_sets(train_path, img_size, classes, validation_size=validation_size)


print("Complete reading input data. Will Now print a snippet of it")
print("Number of files in Training-set:\t\t{}".format(len(data.train.labels)))
print("Number of files in Validation-set:\t{}".format(len(data.valid.labels)))



session = tf.Session()

with tf.name_scope('freeze_layers'):
	x = tf.placeholder(tf.float32, shape=[None, img_size,img_size,num_channels], name='x')

## labels
y_true = tf.placeholder(tf.float32, shape=[None, num_classes], name='y_true')
y_true_cls = tf.argmax(y_true, dimension=1)

def create_weights(shape):
	with tf.name_scope('trainable_layers'):
		x = tf.Variable(tf.truncated_normal(shape, stddev=0.05))
	return x

def create_biases(size):
	with tf.name_scope('trainable_layers'):
		x = tf.Variable(tf.constant(0.05, shape=[size]))
	return x

def conv_layer(input, weights, biases, name):

	layer = tf.nn.conv2d(
		    input=input,
		    filter=weights,
		    strides=[1, 3, 3, 1],
		    padding="SAME",
			name = name)

	layer = tf.nn.bias_add(layer, biases)

	return tf.nn.relu(layer)

def max_pool(input,name):

	pool_layer = tf.nn.max_pool(
			value=input,
			ksize=[1, 2, 2, 1], 	
			strides=[1, 2, 2, 1],
			padding="SAME",
			name=name)

	return pool_layer

def create_flatten_layer(layer):
    #We know that the shape of the layer will be [batch_size img_size img_size num_channels] 
    # But let's get it from the previous layer.
	with tf.name_scope('trainable_layers'):

	    layer_shape = layer.get_shape()

	    ## Number of features will be img_height * img_width* num_channels. But we shall calculate it in place of hard-coding it.
	    num_features = layer_shape[1:4].num_elements()

	    ## Now, we Flatten the layer so we shall have to reshape to num_features
	    layer = tf.reshape(layer, [-1, num_features])

	return layer


def create_fc_layer(input,          
             num_inputs,    
             num_outputs,
             use_relu=True):
    
    #Let's define trainable weights and biases.
    with tf.name_scope('trainable_layers'):
	    weights = create_weights(shape=[num_inputs, num_outputs])
	    biases = create_biases(num_outputs)
	 
	    layer = tf.matmul(input, weights) + biases
	    if use_relu:
	        layer = tf.nn.relu(layer)
	 
    return layer


conv1_1_weights = vgg16_weights['conv1_1_W']
conv1_1_biases = vgg16_weights['conv1_1_b']

conv1_2_weights = vgg16_weights['conv1_2_W']
conv1_2_biases = vgg16_weights['conv1_2_b']

conv2_1_weights = vgg16_weights['conv2_1_W']
conv2_1_biases = vgg16_weights['conv2_1_b']

conv2_2_weights = vgg16_weights['conv2_2_W']
conv2_2_biases = vgg16_weights['conv2_2_b']

conv3_1_weights = vgg16_weights['conv3_1_W']
conv3_1_biases = vgg16_weights['conv3_1_b']

conv3_2_weights = vgg16_weights['conv3_2_W']
conv3_2_biases = vgg16_weights['conv3_2_b']

conv3_3_weights = vgg16_weights['conv3_3_W']
conv3_3_biases = vgg16_weights['conv3_3_b']

conv4_1_weights = vgg16_weights['conv4_1_W']
conv4_1_biases = vgg16_weights['conv4_1_b']

conv4_2_weights = vgg16_weights['conv4_2_W']
conv4_2_biases = vgg16_weights['conv4_2_b']

conv4_3_weights = vgg16_weights['conv4_3_W']
conv4_3_biases = vgg16_weights['conv4_3_b']

conv5_1_weights = vgg16_weights['conv5_1_W']
conv5_1_biases = vgg16_weights['conv5_1_b']

conv5_2_weights = vgg16_weights['conv5_2_W']
conv5_2_biases = vgg16_weights['conv5_2_b']

conv5_3_weights = vgg16_weights['conv5_3_W']
conv5_3_biases = vgg16_weights['conv5_3_b']

with tf.name_scope('trainable_layers'):
	conv1_1 = conv_layer(x,conv1_1_weights,conv1_1_biases,"conv1_1")
	conv1_2 = conv_layer(conv1_1, conv1_2_weights,conv1_2_biases,"conv1_2")
	pool1 = max_pool(conv1_2,"pool1")


	conv2_1 = conv_layer(pool1, conv2_1_weights, conv2_1_biases, "conv2_1")
	conv2_2 = conv_layer(conv2_1, conv2_2_weights, conv2_2_biases, "conv2_2")
	pool2 = max_pool(conv2_2, "pool2")

	conv3_1 = conv_layer(pool2, conv3_1_weights, conv3_1_biases, "conv3_1")
	conv3_2 = conv_layer(conv3_1, conv3_2_weights, conv3_2_biases, "conv3_2")
	conv3_3 = conv_layer(conv3_2, conv3_3_weights,conv3_3_biases, "conv3_3")
	pool3 = max_pool(conv3_3, "pool3")

	conv4_1 = conv_layer(pool3, conv4_1_weights, conv4_1_biases, "conv4_1")
	conv4_2 = conv_layer(conv4_1, conv4_2_weights, conv4_2_biases, "conv4_2")
	conv4_3 = conv_layer(conv4_2, conv4_3_weights, conv4_3_biases, "conv4_3")
	pool4 = max_pool(conv4_3, "pool4")

with tf.name_scope('trainable_layers'):
	conv5_1 = conv_layer(pool4, conv5_1_weights, conv5_1_biases, "conv5_1")
	conv5_2 = conv_layer(conv5_1, conv5_2_weights, conv5_2_biases, "conv5_2")
	conv5_3 = conv_layer(conv5_2, conv5_3_weights, conv5_3_biases, "conv5_3")
	pool5 = max_pool(conv5_3, "pool5")

	flatten_layer = create_flatten_layer(pool5)

	fc_layer1 = create_fc_layer(input=flatten_layer, 
							num_inputs=flatten_layer.get_shape()[1:4].num_elements(),
							num_outputs=fc_layer_size,
							use_relu=True)

	fc_layer2 = create_fc_layer(input=fc_layer1,
	                     num_inputs=fc_layer_size,
	                     num_outputs=num_classes,
	                     use_relu=False)

y_pred = tf.nn.softmax(fc_layer2, name='y_pred')

y_pred_cls = tf.argmax(y_pred, dimension=1)
#session.run(tf.global_variables_initializer())
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=fc_layer2,
                                                    labels=y_pred)
cost = tf.reduce_mean(cross_entropy)

training_layers = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                     "trainable_layers")

print("Trainable_vars: ",training_layers)

optimizer = tf.train.AdadeltaOptimizer().minimize(cost, var_list=training_layers)

correct_prediction = tf.equal(y_pred_cls, y_true_cls)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))


session.run(tf.global_variables_initializer()) 

def show_progress(epoch, feed_dict_train, feed_dict_validate, val_loss):
    acc = session.run(accuracy, feed_dict=feed_dict_train)
    val_acc = session.run(accuracy, feed_dict=feed_dict_validate)
    msg = "Training Epoch {0} --- Training Accuracy: {1:>6.1%}, Validation Accuracy: {2:>6.1%},  Validation Loss: {3:.3f}"
    print(msg.format(epoch + 1, acc, val_acc, val_loss))

total_iterations = 0

saver = tf.train.Saver()

def train(num_iteration):
    global total_iterations
    
    for i in range(total_iterations,
                   total_iterations + num_iteration):

        x_batch, y_true_batch, _, cls_batch = data.train.next_batch(batch_size)
        x_valid_batch, y_valid_batch, _, valid_cls_batch = data.valid.next_batch(batch_size)

        

        feed_dict_tr = {x: x_batch,
                           y_true: y_true_batch}
        feed_dict_val = {x: x_valid_batch,
                              y_true: y_valid_batch}

        #print(feed_dict_tr)
        session.run(optimizer, feed_dict=feed_dict_tr)
        #print(feed_dict_val)
        if i % int(data.train.num_examples/batch_size) == 0: 
            val_loss = session.run(cost, feed_dict=feed_dict_val)
            epoch = int(i / int(data.train.num_examples/batch_size))    
            
            show_progress(epoch, feed_dict_tr, feed_dict_val, val_loss)
    	    writer = tf.summary.FileWriter('./tf_new/', graph=tf.get_default_graph())
    	    writer.add_graph(session.graph)
            saver.save(session, './driverdistraction-model') 


    total_iterations += num_iteration

train(num_iteration=5000)