.

resnet.py

@@ -1,3 +1,8 @@
+#!/usr/bin/env python3
+
+import keras.backend as K
+assert K.image_dim_ordering() == 'th'
+
 import pickle, time
 import sys
 import numpy as np
@@ -6,7 +11,7 @@ from keras.datasets import mnist
 from keras.layers import BatchNormalization
 from keras.layers import Convolution2D, MaxPooling2D
 from keras.layers import Flatten, Reshape
-from keras.layers import Input, merge, Dense, Activation, Dropout
+from keras.layers import Input, merge, Dense, Activation
 from keras.models import Model
 from keras.utils.np_utils import to_categorical
 
@@ -44,15 +49,16 @@ LRprod = 0.1**(1/20.) # will use a tenth of the learning rate after 20 epochs
 
 use_image_generator = True
 
+def prepare(X, y):
+    X = X.reshape(X.shape[0], 1, width, height).astype('float32') / 255
+    # convert class vectors to binary class matrices
+    Y = to_categorical(y_train, nb_classes)
+    return X, Y
+
 # the data, shuffled and split between train and test sets
 (X_train, y_train), (X_test, y_test) = mnist.load_data()
-X_train = X_train.reshape(X_train.shape[0], 1, width, height)
-X_test = X_test.reshape(X_test.shape[0], 1, width, height)
-X_train = X_train.astype('float32') / 255
-X_test = X_test.astype('float32') / 255
-# convert class vectors to binary class matrices
-Y_train = to_categorical(y_train, nb_classes)
-Y_test = to_categorical(y_test, nb_classes)
+X_train, Y_train = prepare(X_train, y_train)
+X_test,  Y_test  = prepare(X_test,  y_test)
 
 if use_image_generator:
     from keras.preprocessing.image import ImageDataGenerator