From 201e52522119527fb2123951c6d762dc0f62641a Mon Sep 17 00:00:00 2001 From: Thibault de Boissiere Date: Mon, 30 Oct 2017 10:55:42 +1100 Subject: [PATCH] Remove complex model Update to keras 2 --- Colorful/src/data/make_dataset.py | 7 +- Colorful/src/model/main.py | 4 - Colorful/src/model/models_colorful.py | 136 ++++---------------------- Colorful/src/model/train_colorful.py | 11 +-- 4 files changed, 27 insertions(+), 131 deletions(-) diff --git a/Colorful/src/data/make_dataset.py b/Colorful/src/data/make_dataset.py index 45d6a7c..fc0eb5a 100644 --- a/Colorful/src/data/make_dataset.py +++ b/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 diff --git a/Colorful/src/model/main.py b/Colorful/src/model/main.py index 20cb721..8b25038 100644 --- a/Colorful/src/model/main.py +++ b/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,7 +42,6 @@ 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 @@ -52,7 +49,6 @@ def launch_eval(**kwargs): 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 diff --git a/Colorful/src/model/models_colorful.py b/Colorful/src/model/models_colorful.py index 1c56521..292bc84 100644 --- a/Colorful/src/model/models_colorful.py +++ b/Colorful/src/model/models_colorful.py @@ -1,10 +1,10 @@ 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): @@ -12,56 +12,43 @@ 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 @@ -69,7 +56,7 @@ def simple_colorful(nb_classes, img_dim, batch_size, model_name="colorful_simple # 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 diff --git a/Colorful/src/model/train_colorful.py b/Colorful/src/model/train_colorful.py index d6dc8a5..50ad448 100644 --- a/Colorful/src/model/train_colorful.py +++ b/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: