fix dssim, add test (#86)

* fix dssim, add test

* add more tests

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))

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)
 

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__])