diff --git a/keras_contrib/backend/theano_backend.py b/keras_contrib/backend/theano_backend.py index 5667aca..2b5adaf 100644 --- a/keras_contrib/backend/theano_backend.py +++ b/keras_contrib/backend/theano_backend.py @@ -164,7 +164,7 @@ def extract_image_patches(X, ksizes, strides, padding='valid', data_format='chan num_channels = xs[-3] patches = images2neibs(X, ksizes, strides, padding) # Theano is sorting by channel - patches = KTH.reshape(patches, (batch, num_channels, KTH.shape(patches)[0] // num_channels, patch_size, patch_size)) + patches = KTH.reshape(patches, (batch, num_channels, num_rows * num_cols, patch_size, patch_size)) patches = KTH.permute_dimensions(patches, (0, 2, 1, 3, 4)) # arrange in a 2d-grid (rows, cols, channels, px, py) patches = KTH.reshape(patches, (batch, num_rows, num_cols, num_channels, patch_size, patch_size)) diff --git a/keras_contrib/losses.py b/keras_contrib/losses.py index 218defd..0fd7308 100644 --- a/keras_contrib/losses.py +++ b/keras_contrib/losses.py @@ -8,6 +8,8 @@ class DSSIMObjective(): """ 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. + Note : In theano, the `kernel_size` must be a factor of the output size. So 3 could + not be the `kernel_size` for an output of 32. # Arguments k1: Parameter of the SSIM (default 0.01) @@ -22,7 +24,7 @@ class DSSIMObjective(): 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.dim_ordering = K.image_data_format() self.backend = KC.backend() def __int_shape(self, x): @@ -36,6 +38,7 @@ class DSSIMObjective(): 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) diff --git a/tests/keras_contrib/losses_test.py b/tests/keras_contrib/losses_test.py index fb66e1c..225317f 100644 --- a/tests/keras_contrib/losses_test.py +++ b/tests/keras_contrib/losses_test.py @@ -1,10 +1,14 @@ import pytest import numpy as np +from numpy.testing import assert_allclose +from keras.layers import Conv2D +from keras.models import Sequential +from keras.optimizers import Adam from keras_contrib import losses from keras import backend as K from keras_contrib import backend as KC - +from keras_contrib.losses import DSSIMObjective allobj = [] @@ -36,5 +40,65 @@ def test_cce_one_hot(): assert K.eval(losses.sparse_categorical_crossentropy(y_a, y_b)).shape == (6,) +def test_DSSIM_channels_last(): + prev_data = K.image_data_format() + K.set_image_data_format('channels_last') + for input_dim, kernel_size in zip([32, 33], [2, 3]): + input_shape = [input_dim, input_dim, 3] + X = np.random.random_sample(4 * input_dim * input_dim * 3).reshape([4] + input_shape) + y = np.random.random_sample(4 * input_dim * input_dim * 3).reshape([4] + input_shape) + + model = Sequential() + model.add(Conv2D(32, (3, 3), padding='same', input_shape=input_shape, activation='relu')) + model.add(Conv2D(3, (3, 3), padding='same', input_shape=input_shape, activation='relu')) + adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-8) + model.compile(loss=DSSIMObjective(kernel_size=kernel_size), metrics=['mse'], optimizer=adam) + model.fit(X, y, batch_size=2, epochs=1, shuffle='batch') + + # Test same + x1 = K.constant(X, 'float32') + x2 = K.constant(X, 'float32') + dssim = DSSIMObjective(kernel_size=kernel_size) + assert_allclose(0.0, K.eval(dssim(x1, x2)), atol=1e-4) + + # Test opposite + x1 = K.zeros([4] + input_shape) + x2 = K.ones([4] + input_shape) + dssim = DSSIMObjective(kernel_size=kernel_size) + assert_allclose(0.5, K.eval(dssim(x1, x2)), atol=1e-4) + + K.set_image_data_format(prev_data) + + +def test_DSSIM_channels_first(): + prev_data = K.image_data_format() + K.set_image_data_format('channels_first') + for input_dim, kernel_size in zip([32, 33], [2, 3]): + input_shape = [3, input_dim, input_dim] + X = np.random.random_sample(4 * input_dim * input_dim * 3).reshape([4] + input_shape) + y = np.random.random_sample(4 * input_dim * input_dim * 3).reshape([4] + input_shape) + + model = Sequential() + model.add(Conv2D(32, (3, 3), padding='same', input_shape=input_shape, activation='relu')) + model.add(Conv2D(3, (3, 3), padding='same', input_shape=input_shape, activation='relu')) + adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-8) + model.compile(loss=DSSIMObjective(kernel_size=kernel_size), metrics=['mse'], optimizer=adam) + model.fit(X, y, batch_size=2, epochs=1, shuffle='batch') + + # Test same + x1 = K.constant(X, 'float32') + x2 = K.constant(X, 'float32') + dssim = DSSIMObjective(kernel_size=kernel_size) + assert_allclose(0.0, K.eval(dssim(x1, x2)), atol=1e-4) + + # Test opposite + x1 = K.zeros([4] + input_shape) + x2 = K.ones([4] + input_shape) + dssim = DSSIMObjective(kernel_size=kernel_size) + assert_allclose(0.5, K.eval(dssim(x1, x2)), atol=1e-4) + + K.set_image_data_format(prev_data) + + if __name__ == '__main__': pytest.main([__file__])