Image Classification with TFLearn

# Residual blocks

# 32 layers: n=5, 56 layers: n=9, 110 layers: n=18

n = 5

batch_size = 1000

num_classes = 10

epochs_shortrun = 5

epochs_longrun = 200

random_seed = 343

save_dir = "/files"

res_dir = "/results"

model_name = 'resnet_cifar10'

import os

save_fn = model_name + ".tfsave"

save_file = os.path.join(save_dir, save_fn)

if not os.path.isdir(res_dir):

  os.makedirs(res_dir)

ckpt_dir = os.path.join(save_dir,"checkpoints")

if not os.path.isdir(ckpt_dir):

    os.makedirs(ckpt_dir)

tblog_dir = os.path.join(save_dir,"tflogs")

if not os.path.isdir(tblog_dir):

    os.makedirs(tblog_dir)

event_dir = os.path.join(tblog_dir,model_name)

exec(open("settings.py").read())

import numpy as np

import dill as pickle

from math import *

def setup_tf(seed):

  import os,random,numpy as np,tensorflow as tf

  from tflearn import init_graph

  tf.reset_default_graph()

  # set random seeds for reproducibility

  os.environ['PYTHONHASHSEED']=str(seed)

  random.seed(seed)

  np.random.seed(seed)

  tf.set_random_seed(seed)

  init_graph(seed=seed)

def setup_cifar_data(verbose=False):

  from tflearn.datasets import cifar10

  from tflearn.data_utils import to_categorical

  datadir = "/cifar/"

  datafile = datadir+"cifar-10-python.tar.gz"

  import shutil,tarfile

  with tarfile.open(datafile, "r:gz") as f:

    f.extractall(datadir)

  # Data loading

  (X, Y), (testX, testY) = cifar10.load_data(dirname="/cifar")

  Y = to_categorical(Y, num_classes)

  testY = to_categorical(testY, num_classes)

  return X, Y, testX, testY

def setup_cifar_labels():

  with open("/cifar/cifar-10-batches-py/batches.meta", 'rb') as fo:

    labels = pickle.load(fo)

  return labels

# Function to find latest checkpoint file

def last_ckpt(dir):

  fl = os.listdir(dir)

  fl = [x for x in fl if x.endswith(".index")]

  cf = ""

  if len(fl) > 0:

    steps = [float(x.split("-")[1][0:-6]) for x in fl]

    m = max(steps)

    cf = fl[steps.index(m)]

    cf = os.path.join(dir,cf)

  return(cf)

def load_model_from_file(model,file):

  # load data from tarfile to model

  import tarfile

  with tarfile.open(file, "r:bz2") as tar:

    try:

      tar.getmember(save_fn+".index")

      tar.getmember(save_fn+".meta")

      tar.getmember(save_fn+".data-00000-of-00001")

    except KeyError:

      print("Minimum training results files not found!\n")

    tar.extractall(path=save_dir)

    print("Loading {}...".format(save_file))

    model.load(save_file, weights_only=False)

def cifar_grid(Xset,Yset,inds,n_col, predictions=None):

  #Visualizing CIFAR 10, takes indicides and shows in a grid

  import matplotlib.pyplot as plt

  if predictions is not None:

    if Yset.shape != predictions.shape:

      print("Predictions must equal Yset in length!")

      return(None)

  N = len(inds)

  n_row = int(ceil(1.0*N/n_col))

  fig, axes = plt.subplots(n_row,n_col,figsize=(10,10))

  clabels = labels["label_names"]

  for j in range(n_row):

    for k in range(n_col):

      i_inds = j*n_col+k

      i_data = inds[i_inds]

      axes[j][k].set_axis_off()

      if i_inds < N:

        axes[j][k].imshow(Xset[i_data,...], interpolation='nearest')

        label = clabels[np.argmax(Yset[i_data,...])]

        axes[j][k].set_title(label)

        if predictions is not None:

          pred = clabels[np.argmax(predictions[i_data,...])]

          if label != pred:

            label += " n"

            axes[j][k].set_title(pred, color='red')            

  fig.set_tight_layout(True)

  return fig

from tflearn import ImagePreprocessing, ImageAugmentation

from tflearn import input_data, DNN

from tflearn import conv_2d, residual_block

from tflearn import batch_normalization, activation, global_avg_pool 

from tflearn import fully_connected, Momentum, regression

from tflearn.callbacks import Callback

print("Using real-time data augmentation.\n")

# Real-time data preprocessing

img_prep = ImagePreprocessing()

img_prep.add_featurewise_zero_center(per_channel=True, 

  mean=[ 0.49139968, 0.48215841, 0.44653091 ])

# Real-time data augmentation

img_aug = ImageAugmentation()

img_aug.add_random_flip_leftright()

img_aug.add_random_crop([32, 32], padding=4)

# Building Residual Network

net = input_data(shape=[None, 32, 32, 3],

                 data_preprocessing=img_prep,

                 data_augmentation=img_aug)

net = conv_2d(net, 16, 3, regularizer='L2', weight_decay=0.0001)

net = residual_block(net, n, 16)

net = residual_block(net, 1, 32, downsample=True)

net = residual_block(net, n-1, 32)

net = residual_block(net, 1, 64, downsample=True)

net = residual_block(net, n-1, 64)

net = batch_normalization(net)

net = activation(net, 'relu')

net = global_avg_pool(net)

# Regression

net = fully_connected(net, num_classes, activation='softmax')

mom = Momentum(0.1, lr_decay=0.1, decay_step=32000, staircase=True)

net = regression(net, optimizer=mom, loss='categorical_crossentropy')

# Initialize model

ckpt_file = os.path.join(ckpt_dir,"model.ckpt")

model = DNN(net, checkpoint_path=ckpt_file,

                    max_checkpoints=10, clip_gradients=0.,

                    tensorboard_dir=tblog_dir,tensorboard_verbose=0)

# disabled until directories can be written to /results

#cff = last_ckpt(ckpt_dir)

#if cff != "":

#  print("Loading ",cff,"...")

#  model.load(cff)

# define the early-stop callback

class EarlyStoppingCallback(Callback):

    def __init__(self, val_loss_thresh, val_loss_patience):

        """ minimum loss improvement setup """

        self.val_loss_thresh = val_loss_thresh

        self.val_loss_last = float('inf')

        self.val_loss_patience = val_loss_patience

        self.val_loss_squint = 0

    def on_batch_end(self, training_state, snapshot=False):

        """ loss improvement threshold w/ patience """

        # Apparently this can happen.

        if training_state.val_loss is None: return

        if (self.val_loss_last 

          - training_state.val_loss) < self.val_loss_thresh:

          # unacceptable!

          if self.val_loss_squint >= self.val_loss_patience:

            raise StopIteration

          else:

            self.val_loss_squint += 1

        else:

          # we good again - reset

          self.val_loss_last = training_state.val_loss

          self.val_loss_squint = 0

# Initialize our callback.

early_stopping_cb = EarlyStoppingCallback(

  val_loss_thresh=0.001, 

  val_loss_patience=25

)