231 lines
10 KiB
Python
231 lines
10 KiB
Python
|
import numpy as np
|
||
|
|
import pandas as pd
|
||
|
|
from time import time
|
||
|
|
from sklearn.model_selection import train_test_split
|
||
|
|
import keras
|
||
|
|
import keras.layers as layers
|
||
|
|
from keras import optimizers
|
||
|
|
from keras.preprocessing.image import ImageDataGenerator
|
||
|
|
from keras.models import Model
|
||
|
|
from keras.layers import Input, Activation, Concatenate
|
||
|
|
from keras.layers import Flatten, Dropout
|
||
|
|
from keras.layers import Convolution2D, MaxPooling2D
|
||
|
|
from keras.layers import GlobalAveragePooling2D
|
||
|
|
|
||
|
|
# Training variables
|
||
|
|
EPOCHS = 30
|
||
|
|
BATCH_SIZE = 480
|
||
|
|
|
||
|
|
# dataset
|
||
|
|
def prepare_mnist_data(rows=28, cols=28, nb_classes=10, categorical=False, padding = True, debug = True):
|
||
|
|
""" Get MNIST data """
|
||
|
|
|
||
|
|
from keras.datasets import mnist
|
||
|
|
from keras.utils import np_utils
|
||
|
|
(X_train, y_train), (X_test, y_test) = mnist.load_data()
|
||
|
|
X_train = X_train.reshape(X_train.shape[0], 1, rows, cols)
|
||
|
|
X_test = X_test.reshape(X_test.shape[0], 1, rows, cols)
|
||
|
|
X_train = X_train.astype('float32')
|
||
|
|
X_test = X_test.astype('float32')
|
||
|
|
X_train /= 255
|
||
|
|
X_test /= 255
|
||
|
|
X_train, X_validation, y_train, y_validation = train_test_split(X_train, y_train, test_size=0.2, random_state=0)
|
||
|
|
if padding:
|
||
|
|
# Pad images with 0s
|
||
|
|
X_train = np.pad(X_train, ((0,0),(0,0),(2,2),(2,2)), 'constant')
|
||
|
|
X_test = np.pad(X_test, ((0,0),(0,0),(2,2),(2,2)), 'constant')
|
||
|
|
X_validation = np.pad(X_validation, ((0,0),(0,0),(2,2),(2,2)), 'constant')
|
||
|
|
|
||
|
|
if categorical:
|
||
|
|
# convert class vectors to binary class matrices
|
||
|
|
y_train = np_utils.to_categorical(y_train, nb_classes)
|
||
|
|
y_test = np_utils.to_categorical(y_test, nb_classes)
|
||
|
|
y_validation = np_utils.to_categorical(y_validation, nb_classes)
|
||
|
|
|
||
|
|
if debug:
|
||
|
|
print('X_train shape:', X_train.shape)
|
||
|
|
print(X_train.shape[0], 'train samples')
|
||
|
|
print(X_test.shape[0], 'test samples')
|
||
|
|
print(X_validation.shape[0], 'validation samples')
|
||
|
|
if not categorical:
|
||
|
|
train_labels_count = np.unique(y_train, return_counts=True)
|
||
|
|
dataframe_train_labels = pd.DataFrame({'Label':train_labels_count[0], 'Count':train_labels_count[1]})
|
||
|
|
print(dataframe_train_labels)
|
||
|
|
|
||
|
|
return X_train, X_test, X_validation, y_train, y_test, y_validation
|
||
|
|
|
||
|
|
|
||
|
|
# SqueezeNet
|
||
|
|
def SqueezeNet(nb_classes, inputs=(1, 32, 32)):
|
||
|
|
""" Keras Implementation of SqueezeNet(arXiv 1602.07360)
|
||
|
|
@param nb_classes: total number of final categories
|
||
|
|
Arguments:
|
||
|
|
inputs -- shape of the input images (channel, cols, rows)
|
||
|
|
"""
|
||
|
|
|
||
|
|
input_img = Input(shape=inputs)
|
||
|
|
conv1 = Convolution2D(
|
||
|
|
96, (7, 7), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
strides=(2, 2), padding='same', name='conv1',
|
||
|
|
data_format="channels_first")(input_img)
|
||
|
|
maxpool1 = MaxPooling2D(
|
||
|
|
pool_size=(3, 3), strides=(2, 2), name='maxpool1',
|
||
|
|
data_format="channels_first")(conv1)
|
||
|
|
fire2_squeeze = Convolution2D(
|
||
|
|
16, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire2_squeeze',
|
||
|
|
data_format="channels_first")(maxpool1)
|
||
|
|
fire2_expand1 = Convolution2D(
|
||
|
|
64, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire2_expand1',
|
||
|
|
data_format="channels_first")(fire2_squeeze)
|
||
|
|
fire2_expand2 = Convolution2D(
|
||
|
|
64, (3, 3), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire2_expand2',
|
||
|
|
data_format="channels_first")(fire2_squeeze)
|
||
|
|
merge2 = Concatenate(axis=1)([fire2_expand1, fire2_expand2])
|
||
|
|
|
||
|
|
fire3_squeeze = Convolution2D(
|
||
|
|
16, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire3_squeeze',
|
||
|
|
data_format="channels_first")(merge2)
|
||
|
|
fire3_expand1 = Convolution2D(
|
||
|
|
64, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire3_expand1',
|
||
|
|
data_format="channels_first")(fire3_squeeze)
|
||
|
|
fire3_expand2 = Convolution2D(
|
||
|
|
64, (3, 3), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire3_expand2',
|
||
|
|
data_format="channels_first")(fire3_squeeze)
|
||
|
|
merge3 = Concatenate(axis=1)([fire3_expand1, fire3_expand2])
|
||
|
|
|
||
|
|
fire4_squeeze = Convolution2D(
|
||
|
|
32, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire4_squeeze',
|
||
|
|
data_format="channels_first")(merge3)
|
||
|
|
fire4_expand1 = Convolution2D(
|
||
|
|
128, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire4_expand1',
|
||
|
|
data_format="channels_first")(fire4_squeeze)
|
||
|
|
fire4_expand2 = Convolution2D(
|
||
|
|
128, (3, 3), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire4_expand2',
|
||
|
|
data_format="channels_first")(fire4_squeeze)
|
||
|
|
merge4 = Concatenate(axis=1)([fire4_expand1, fire4_expand2])
|
||
|
|
maxpool4 = MaxPooling2D(
|
||
|
|
pool_size=(3, 3), strides=(2, 2), name='maxpool4',
|
||
|
|
data_format="channels_first")(merge4)
|
||
|
|
|
||
|
|
fire5_squeeze = Convolution2D(
|
||
|
|
32, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire5_squeeze',
|
||
|
|
data_format="channels_first")(maxpool4)
|
||
|
|
fire5_expand1 = Convolution2D(
|
||
|
|
128, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire5_expand1',
|
||
|
|
data_format="channels_first")(fire5_squeeze)
|
||
|
|
fire5_expand2 = Convolution2D(
|
||
|
|
128, (3, 3), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire5_expand2',
|
||
|
|
data_format="channels_first")(fire5_squeeze)
|
||
|
|
merge5 = Concatenate(axis=1)([fire5_expand1, fire5_expand2])
|
||
|
|
|
||
|
|
fire6_squeeze = Convolution2D(
|
||
|
|
48, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire6_squeeze',
|
||
|
|
data_format="channels_first")(merge5)
|
||
|
|
fire6_expand1 = Convolution2D(
|
||
|
|
192, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire6_expand1',
|
||
|
|
data_format="channels_first")(fire6_squeeze)
|
||
|
|
fire6_expand2 = Convolution2D(
|
||
|
|
192, (3, 3), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire6_expand2',
|
||
|
|
data_format="channels_first")(fire6_squeeze)
|
||
|
|
merge6 = Concatenate(axis=1)([fire6_expand1, fire6_expand2])
|
||
|
|
|
||
|
|
fire7_squeeze = Convolution2D(
|
||
|
|
48, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire7_squeeze',
|
||
|
|
data_format="channels_first")(merge6)
|
||
|
|
fire7_expand1 = Convolution2D(
|
||
|
|
192, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire7_expand1',
|
||
|
|
data_format="channels_first")(fire7_squeeze)
|
||
|
|
fire7_expand2 = Convolution2D(
|
||
|
|
192, (3, 3), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire7_expand2',
|
||
|
|
data_format="channels_first")(fire7_squeeze)
|
||
|
|
merge7 = Concatenate(axis=1)([fire7_expand1, fire7_expand2])
|
||
|
|
|
||
|
|
fire8_squeeze = Convolution2D(
|
||
|
|
64, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire8_squeeze',
|
||
|
|
data_format="channels_first")(merge7)
|
||
|
|
fire8_expand1 = Convolution2D(
|
||
|
|
256, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire8_expand1',
|
||
|
|
data_format="channels_first")(fire8_squeeze)
|
||
|
|
fire8_expand2 = Convolution2D(
|
||
|
|
256, (3, 3), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire8_expand2',
|
||
|
|
data_format="channels_first")(fire8_squeeze)
|
||
|
|
merge8 = Concatenate(axis=1)([fire8_expand1, fire8_expand2])
|
||
|
|
|
||
|
|
maxpool8 = MaxPooling2D(
|
||
|
|
pool_size=(3, 3), strides=(2, 2), name='maxpool8',
|
||
|
|
data_format="channels_first")(merge8)
|
||
|
|
fire9_squeeze = Convolution2D(
|
||
|
|
64, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire9_squeeze',
|
||
|
|
data_format="channels_first")(maxpool8)
|
||
|
|
fire9_expand1 = Convolution2D(
|
||
|
|
256, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire9_expand1',
|
||
|
|
data_format="channels_first")(fire9_squeeze)
|
||
|
|
fire9_expand2 = Convolution2D(
|
||
|
|
256, (3, 3), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='same', name='fire9_expand2',
|
||
|
|
data_format="channels_first")(fire9_squeeze)
|
||
|
|
merge9 = Concatenate(axis=1)([fire9_expand1, fire9_expand2])
|
||
|
|
|
||
|
|
fire9_dropout = Dropout(0.5, name='fire9_dropout')(merge9)
|
||
|
|
conv10 = Convolution2D(
|
||
|
|
nb_classes, (1, 1), activation='relu', kernel_initializer='glorot_uniform',
|
||
|
|
padding='valid', name='conv10',
|
||
|
|
data_format="channels_first")(fire9_dropout)
|
||
|
|
|
||
|
|
global_avgpool10 = GlobalAveragePooling2D(data_format='channels_first')(conv10)
|
||
|
|
softmax = Activation("softmax", name='softmax')(global_avgpool10)
|
||
|
|
|
||
|
|
return Model(inputs=input_img, outputs=softmax)
|
||
|
|
|
||
|
|
|
||
|
|
# Get dataset
|
||
|
|
X_train, X_test, X_validation, y_train, y_test, y_validation = prepare_mnist_data(categorical = True)
|
||
|
|
|
||
|
|
# Check
|
||
|
|
model = SqueezeNet(10)
|
||
|
|
model.summary()
|
||
|
|
|
||
|
|
# Preparation
|
||
|
|
model.compile(loss=keras.losses.categorical_crossentropy, optimizer=keras.optimizers.Adam(), metrics=['accuracy'])
|
||
|
|
|
||
|
|
# Trainer and validator
|
||
|
|
train_generator = ImageDataGenerator().flow(X_train, y_train, batch_size=BATCH_SIZE)
|
||
|
|
validation_generator = ImageDataGenerator().flow(X_validation, y_validation, batch_size=BATCH_SIZE)
|
||
|
|
|
||
|
|
# Training
|
||
|
|
from keras.callbacks import TensorBoard
|
||
|
|
steps_per_epoch = X_train.shape[0]//BATCH_SIZE
|
||
|
|
validation_steps = X_validation.shape[0]//BATCH_SIZE
|
||
|
|
tensorboard = TensorBoard(log_dir="logs/{}".format(time()))
|
||
|
|
|
||
|
|
model.fit_generator(train_generator, steps_per_epoch=steps_per_epoch, epochs=EPOCHS,
|
||
|
|
validation_data=validation_generator, validation_steps=validation_steps,
|
||
|
|
shuffle=True, callbacks=[tensorboard])
|
||
|
|
|
||
|
|
# Saving the model in keras format (.h5)
|
||
|
|
model.save('squeezenet_keras_v2.2.4.h5')
|
||
|
|
print(">>> Model saved!")
|