Remove complex model Update to keras 2

Colorful/src/data/make_dataset.py

@@ -4,7 +4,10 @@
 import parmap
 import argparse
 import numpy as np
-import cPickle as pickle
+try:
+    import cPickle as pickle
+except ImportError:
+    import pickle
 from skimage import color
 from tqdm import tqdm as tqdm
 import sklearn.neighbors as nn
@@ -52,7 +55,7 @@ def build_HDF5(size=64):
             img = celeb[0]
             attrs = int(celeb[1])
             d_partition[img] = attrs
-    with open(os.path.join(data_dir, "d_partition.pickle"), "w") as fd:
+    with open(os.path.join(data_dir, "d_partition.pickle"), "wb") as fd:
         pickle.dump(d_partition, fd)
 
     # Put train data in HDF5

Colorful/src/model/main.py

@@ -20,8 +20,6 @@ def launch_eval(**kwargs):
     parser = argparse.ArgumentParser(description='Train model')
     parser.add_argument('mode', type=str, help="Choose train or eval")
     parser.add_argument('data_file', type=str, help="Path to HDF5 containing the data")
-    parser.add_argument('--model_name', type=str, default="simple_colorful",
-                        help="Model name. Choose simple_colorful or colorful")
     parser.add_argument('--training_mode', default="in_memory", type=str,
                         help=('Training mode. Choose in_memory to load all the data in memory and train.'
                               'Choose on_demand to load batches from disk at each step'))
@@ -44,15 +42,13 @@ def launch_eval(**kwargs):
                 "nb_epoch": args.nb_epoch,
                 "nb_resblocks":args.nb_resblocks,
                 "training_mode": args.training_mode,
-                "model_name": args.model_name,
                 "nb_neighbors": args.nb_neighbors,
                 "epoch": args.epoch,
                 "T": args.T
                 }
 
     assert args.mode in ["train", "eval"]
     assert args.training_mode in ["in_memory", "on_demand"]
-    assert args.model_name in ["colorful", "simple_colorful"]
 
     if args.mode == "train":
         # Launch training

Colorful/src/model/models_colorful.py

@@ -1,75 +1,62 @@
 import keras.backend as K
 from keras.models import Model
 from keras.regularizers import l2
-from keras.layers import Input, merge
+from keras.layers import Input, Add
 from keras.layers.core import Activation, Lambda
 from keras.layers.normalization import BatchNormalization
-from keras.layers.convolutional import Convolution2D, AtrousConvolution2D, UpSampling2D
+from keras.layers.convolutional import Conv2D, UpSampling2D
 
 
 def residual_block(x, nb_filter, block_idx, bn=True, weight_decay=0):
 
     # 1st conv
     name = "block%s_conv2D%s" % (block_idx, "a")
     W_reg = l2(weight_decay)
-    r = Convolution2D(nb_filter, 3, 3, border_mode="same", W_regularizer=W_reg, name=name)(x)
+    r = Conv2D(nb_filter, (3, 3), padding="same", kernel_regularizer=W_reg, name=name)(x)
     if bn:
-        r = BatchNormalization(mode=2, axis=1, name="block%s_bn%s" % (block_idx, "a"))(r)
+        r = BatchNormalization(axis=1, name="block%s_bn%s" % (block_idx, "a"))(r)
     r = Activation("relu", name="block%s_relu%s" % (block_idx, "a"))(r)
 
     # 2nd conv
     name = "block%s_conv2D%s" % (block_idx, "b")
     W_reg = l2(weight_decay)
-    r = Convolution2D(nb_filter, 3, 3, border_mode="same", W_regularizer=W_reg, name=name)(r)
+    r = Conv2D(nb_filter, (3, 3), padding="same", kernel_regularizer=W_reg, name=name)(r)
     if bn:
-        r = BatchNormalization(mode=2, axis=1, name="block%s_bn%s" % (block_idx, "b"))(r)
+        r = BatchNormalization(axis=1, name="block%s_bn%s" % (block_idx, "b"))(r)
     r = Activation("relu", name="block%s_relu%s" % (block_idx, "b"))(r)
 
     # Merge residual and identity
-    x = merge([x, r], mode='sum', concat_axis=1, name="block%s_merge" % block_idx)
+    x = Add(name="block%s_merge" % block_idx)([x, r])
 
     return x
 
 
-def convolutional_block(x, block_idx, nb_filter, nb_conv, subsample):
+def convolutional_block(x, block_idx, nb_filter, nb_conv, strides):
 
     # 1st conv
     for i in range(nb_conv):
         name = "block%s_conv2D_%s" % (block_idx, i)
         if i < nb_conv - 1:
-            x = Convolution2D(nb_filter, 3, 3, name=name, border_mode="same")(x)
-            x = BatchNormalization(mode=2, axis=1)(x)
+            x = Conv2D(nb_filter, (3, 3), name=name, padding="same")(x)
+            x = BatchNormalization(axis=1)(x)
             x = Activation("relu")(x)
         else:
-            x = Convolution2D(nb_filter, 3, 3, name=name, subsample=subsample, border_mode="same")(x)
-            x = BatchNormalization(mode=2, axis=1)(x)
+            x = Conv2D(nb_filter, (3, 3), name=name, strides=strides, padding="same")(x)
+            x = BatchNormalization(axis=1)(x)
             x = Activation("relu")(x)
 
     return x
 
 
-def atrous_block(x, block_idx, nb_filter, nb_conv):
-
-    # 1st conv
-    for i in range(nb_conv):
-        name = "block%s_conv2D_%s" % (block_idx, i)
-        x = AtrousConvolution2D(nb_filter, 3, 3, name=name, border_mode="same")(x)
-        x = BatchNormalization(mode=2, axis=1)(x)
-
-        x = Activation("relu")(x)
-
-    return x
-
-
-def simple_colorful(nb_classes, img_dim, batch_size, model_name="colorful_simple"):
+def colorful(nb_classes, img_dim, batch_size, model_name="colorful"):
 
     nb_resblocks = 3
     block_idx = 0
     h, w = img_dim[1:]
 
     # First conv block
     x_input = Input(shape=img_dim, name="input")
-    x = Convolution2D(64, 3, 3, name="block%s_conv2d_0" % block_idx, border_mode="same")(x_input)
+    x = Conv2D(64, (3, 3), name="block%s_conv2d_0" % block_idx, padding="same")(x_input)
     x = Activation("relu", name="block%s_relu" % block_idx)(x)
     block_idx += 1
 
@@ -79,7 +66,7 @@ def simple_colorful(nb_classes, img_dim, batch_size, model_name="colorful_simple
         block_idx += 1
 
     # Final conv
-    x = Convolution2D(nb_classes, 1, 1, name="final_conv2D", border_mode="same")(x)
+    x = Conv2D(nb_classes, (1, 1), name="final_conv2D", padding="same")(x)
 
     # Reshape Softmax
     def output_shape(input_shape):
@@ -99,103 +86,14 @@ def reshape_softmax(x):
     ReshapeSoftmax = Lambda(lambda z: reshape_softmax(z), output_shape=output_shape, name="ReshapeSoftmax")
     x = ReshapeSoftmax(x)
 
-    # Build model
-    colorful_simple = Model(input=[x_input], output=[x], name=model_name)
-
-    return colorful_simple
-
-
-def colorful(nb_classes, img_dim, batch_size, model_name="colorful"):
-    """
-    """
-    h, w = img_dim[1:]
-
-    x_input = Input(shape=img_dim, name="input")
-
-    # Keep track of image h and w
-    current_h, current_w = img_dim[1:]
-
-    # Convolutional blocks parameters
-    list_filter_size = [64, 128, 256, 512, 512]
-    list_block_size = [2, 2, 3, 3, 3]
-    subsample = [(2,2), (2,2), (2,2), (1,1), (1,1)]
-
-    # A trous blocks parameters
-    list_filter_size_atrous = [512, 512]
-    list_block_size_atrous = [3, 3]
-
-    block_idx = 0
-
-    # First block
-    f, b, s = list_filter_size[0], list_block_size[0], subsample[0]
-    x = convolutional_block(x_input, block_idx, f, b, s)
-    block_idx += 1
-    current_h, current_w = current_h / s[0], current_w / s[1]
-
-    # Next blocks
-    for f, b, s in zip(list_filter_size[1:-1], list_block_size[1:-1], subsample[1:-1]):
-        x = convolutional_block(x, block_idx, f, b, s)
-        block_idx += 1
-        current_h, current_w = current_h / s[0], current_w / s[1]
-
-    # Atrous blocks
-    for idx, (f, b) in enumerate(zip(list_filter_size_atrous, list_block_size_atrous)):
-        x = atrous_block(x, block_idx, f, b)
-        block_idx += 1
-
-    # Block 7
-    f, b, s = list_filter_size[-1], list_block_size[-1], subsample[-1]
-    x = convolutional_block(x, block_idx, f, b, s)
-    block_idx += 1
-    current_h, current_w = current_h / s[0], current_w / s[1]
-
-    # Block 8
-    # Not using Deconvolution at the moment
-    # x = Deconvolution2D(256, 2, 2,
-    #                     output_shape=(None, 256, current_h * 2, current_w * 2),
-    #                     subsample=(2, 2),
-    #                     border_mode="valid")(x)
-    x = UpSampling2D(size=(2,2), name="upsampling2d")(x)
-
-    x = convolutional_block(x, block_idx, 256, 2, (1, 1))
-    block_idx += 1
-    current_h, current_w = current_h * 2, current_w * 2
-
-    # Final conv
-    x = Convolution2D(nb_classes, 1, 1, name="conv2d_final", border_mode="same")(x)
-
-    # Reshape Softmax
-    def output_shape(input_shape):
-        return (batch_size, h, w, nb_classes + 1)
-
-    def reshape_softmax(x):
-        x = K.permute_dimensions(x, [0, 2, 3, 1])  # last dimension in number of filters
-        x = K.reshape(x, (batch_size * current_h * current_w, nb_classes))
-        x = K.softmax(x)
-        # Add a zero column so that x has the same dimension as the target (313 classes + 1 weight)
-        xc = K.zeros((batch_size * current_h * current_w, 1))
-        x = K.concatenate([x, xc], axis=1)
-        # Reshape back to (batch_size, h, w, nb_classes + 1) to satisfy keras' shape checks
-        x = K.reshape(x, (batch_size, current_h, current_w, nb_classes + 1))
-        x = K.resize_images(x, h / current_h, w / current_w, "tf")
-        # x = K.permute_dimensions(x, [0, 3, 1, 2])
-        return x
-
-    ReshapeSoftmax = Lambda(lambda z: reshape_softmax(z), output_shape=output_shape, name="ReshapeSoftmax")
-    x = ReshapeSoftmax(x)
-
     # Build model
     colorful = Model(input=[x_input], output=[x], name=model_name)
 
     return colorful
 
 
-def load(model_name, nb_classes, img_dim, batch_size):
-
-    if model_name == "colorful":
-        model = colorful(nb_classes, img_dim, batch_size, model_name=model_name)
+def load(nb_classes, img_dim, batch_size):
 
-    if model_name == "simple_colorful":
-        model = simple_colorful(nb_classes, img_dim, batch_size, model_name=model_name)
+    model = colorful(nb_classes, img_dim, batch_size, model_name="colorful")
 
     return model

Colorful/src/model/train_colorful.py

@@ -49,13 +49,12 @@ def train(**kwargs):
     nb_epoch = kwargs["nb_epoch"]
     data_file = kwargs["data_file"]
     nb_neighbors = kwargs["nb_neighbors"]
-    model_name = kwargs["model_name"]
     training_mode = kwargs["training_mode"]
     epoch_size = n_batch_per_epoch * batch_size
     img_size = int(os.path.basename(data_file).split("_")[1])
 
     # Setup directories to save model, architecture etc
-    general_utils.setup_logging(model_name)
+    general_utils.setup_logging("colorful")
 
     # Create a batch generator for the color data
     DataGen = batch_utils.DataGenerator(data_file,
@@ -87,12 +86,12 @@ def train(**kwargs):
         opt = Adam(lr=1E-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08)
 
         # Load colorizer model
-        color_model = models.load(model_name, nb_q, (1, h, w), batch_size)
+        color_model = models.load(nb_q, (1, h, w), batch_size)
         color_model.compile(loss=categorical_crossentropy_color, optimizer=opt)
 
         color_model.summary()
-        from keras.utils.visualize_util import plot
-        plot(color_model, to_file='../../figures/colorful.png', show_shapes=True, show_layer_names=True)
+        from keras.utils import plot_model
+        plot_model(color_model, to_file='../../figures/colorful.png', show_shapes=True, show_layer_names=True)
 
         # Actual training loop
         for epoch in range(nb_epoch):
@@ -130,7 +129,7 @@ def train(**kwargs):
             # Save weights every 5 epoch
             if epoch % 5 == 0:
                 weights_path = os.path.join('../../models/%s/%s_weights_epoch%s.h5' %
-                                            (model_name, model_name, epoch))
+                                            ("colorful", "colorful", epoch))
                 color_model.save_weights(weights_path, overwrite=True)
 
     except KeyboardInterrupt: