diff --git a/keras_contrib/backend/tensorflow_backend.py b/keras_contrib/backend/tensorflow_backend.py index b70b81a..26d888c 100644 --- a/keras_contrib/backend/tensorflow_backend.py +++ b/keras_contrib/backend/tensorflow_backend.py @@ -2,6 +2,7 @@ import tensorflow as tf from tensorflow.python.training import moving_averages from tensorflow.python.ops import tensor_array_ops from tensorflow.python.ops import control_flow_ops + try: from tensorflow.python.ops import ctc_ops as ctc except ImportError: @@ -28,7 +29,7 @@ def _preprocess_deconv_output_shape(x, shape, dim_ordering): shape = (shape[0],) + tuple(shape[2:]) + (shape[1],) if shape[0] is None: - shape = (tf.shape(x)[0], ) + tuple(shape[1:]) + shape = (tf.shape(x)[0],) + tuple(shape[1:]) shape = tf.stack(list(shape)) return shape @@ -100,9 +101,8 @@ def extract_image_patches(x, ksizes, ssizes, border_mode="same", padding) # Reshaping to fit Theano bs, w, h, ch = KTF.int_shape(patches) - patches = tf.reshape(patches, [bs, w, h, -1, ch_i]) - patches = tf.reshape(tf.transpose(patches, [0, 1, 2, 4, 3]), - [bs, w, h, ch_i, ksizes[0], ksizes[1]]) + patches = tf.reshape(tf.transpose(tf.reshape(patches, [-1, w, h, tf.floordiv(ch, ch_i), ch_i]), [0, 1, 2, 4, 3]), + [-1, w, h, ch_i, ksizes[0], ksizes[1]]) if dim_ordering == "tf": patches = KTF.permute_dimensions(patches, [0, 1, 2, 4, 5, 3]) return patches diff --git a/keras_contrib/objectives.py b/keras_contrib/objectives.py index e40fe99..c3603ed 100644 --- a/keras_contrib/objectives.py +++ b/keras_contrib/objectives.py @@ -1,10 +1,64 @@ from __future__ import absolute_import -from . import backend as K -from keras.utils.generic_utils import get_from_module - from keras.objectives import * +import keras_contrib.backend as KC + def get(identifier): return get_from_module(identifier, globals(), 'objective') + + +class DSSIMObjective(): + def __init__(self, k1=0.01, k2=0.03, kernel_size=3, max_value=1.0): + """ + Difference of Structural Similarity (DSSIM loss function). Clipped between 0 and 0.5 + Note : You should add a regularization term like a l2 loss in addition to this one. + :param batch_size: Batch size used in the model + :param k1: Parameter of the SSIM (default 0.01) + :param k2: Parameter of the SSIM (default 0.03) + :param kernel_size: Size of the sliding window (default 3) + :param max_value: Max value of the output (default 1.0) + """ + self.__name__ = "DSSIMObjective" + self.kernel_size = kernel_size + self.k1 = k1 + self.k2 = k2 + self.max_value = max_value + self.c1 = (self.k1 * self.max_value) ** 2 + self.c2 = (self.k2 * self.max_value) ** 2 + self.dim_ordering = K.image_dim_ordering() + self.backend = KC.backend() + + def __int_shape(self, x): + return KC.int_shape(x) if self.backend == "tensorflow" else KC.shape(x) + + def __call__(self, y_true, y_pred): + # There are additional parameters for this function + # Note: some of the 'modes' for edge behavior do not yet have a gradient definition in the Theano tree + # and cannot be used for learning + + kernel = [self.kernel_size, self.kernel_size] + y_true = KC.reshape(y_true, [-1] + list(self.__int_shape(y_pred)[1:])) + y_pred = KC.reshape(y_pred, [-1] + list(self.__int_shape(y_pred)[1:])) + patches_pred = KC.extract_image_patches(y_pred, kernel, kernel, "valid", self.dim_ordering) + patches_true = KC.extract_image_patches(y_true, kernel, kernel, "valid", self.dim_ordering) + + # Reshape to get the var in the cells + bs, w, h, c1, c2, c3 = self.__int_shape(patches_pred) + patches_pred = KC.reshape(patches_pred, [-1, w, h, c1 * c2 * c3]) + patches_true = KC.reshape(patches_true, [-1, w, h, c1 * c2 * c3]) + # Get mean + u_true = KC.mean(patches_true, axis=-1) + u_pred = KC.mean(patches_pred, axis=-1) + # Get variance + var_true = K.var(patches_true, axis=-1) + var_pred = K.var(patches_pred, axis=-1) + # Get std dev + std_true = K.sqrt(var_true + KC.epsilon()) + std_pred = K.sqrt(var_pred + KC.epsilon()) + + ssim = (2 * u_true * u_pred + self.c1) * (2 * std_pred * std_true + self.c2) + denom = (K.square(u_true) + K.square(u_pred) + self.c1) * (var_pred + var_true + self.c2) + ssim /= denom # no need for clipping, c1 and c2 make the denom non-zero + return K.mean((1.0 - ssim) / 2.0) diff --git a/tests/keras_contrib/test_objectives.py b/tests/keras_contrib/test_objectives.py index ee07515..a50e1d7 100644 --- a/tests/keras_contrib/test_objectives.py +++ b/tests/keras_contrib/test_objectives.py @@ -1,11 +1,11 @@ -import pytest import numpy as np - +import pytest from keras import backend as K +from numpy.testing import assert_allclose + from keras_contrib import backend as KC from keras_contrib import objectives - allobj = [] @@ -25,5 +25,32 @@ def test_objective_shapes_2d(): assert K.eval(objective_output).shape == (6,) +def test_dssim_same(): + x = np.random.random_sample(30 * 30 * 3).reshape([1, 30, 30, 3]) + x1 = KC.variable(x) + loss = objectives.DSSIMObjective() + assert_allclose([0.0], KC.eval(loss(x1, x1)), atol=1.0e-4) + + +def test_dssim_opposite(): + x = np.zeros([1, 30, 30, 3]) + x1 = KC.variable(x) + y = np.ones([1, 30, 30, 3]) + y1 = KC.variable(y) + loss = objectives.DSSIMObjective() + assert_allclose([0.5], KC.eval(loss(x1, y1)), atol=1.0e-4) + + +def test_dssim_compile(): + from keras.models import Sequential + from keras.layers import Convolution2D + x = np.zeros([1, 30, 30, 3]) + loss = objectives.DSSIMObjective() + model = Sequential() + model.add(Convolution2D(3, 3, 3, border_mode="same", input_shape=(30, 30, 3))) + model.compile("rmsprop", loss) + model.fit([x], [x], 1, 1) + + if __name__ == "__main__": pytest.main([__file__])