Keras - MS-SSIM as loss function

Question

Edit: Updates since I had initially misinterpreted the paper

I am trying to implement a custom loss function for keras, such that the objective is to minimize the MS-SSIM (http://www.cns.nyu.edu/~zwang/files/papers/msssim.pdf)

I am getting the following error:

Traceback (most recent call last):
  File "kerasmodel_const_init_customloss.py", line 318, in <module>
    model.fit(x=[np.array(training_data_LR), np.array(training_data_MC)], y=[np.array(training_data_HR)], batch_size=128, epochs=2, verbose=1, validation_data=([np.array(validation_data_LR), np.array(validation_data_MC)], np.array(validation_data_HR)), shuffle=True, callbacks=[log_callback, checkpoint_callback])
  File "/usr/local/lib/python2.7/dist-packages/keras/models.py", line 965, in fit
    validation_steps=validation_steps)
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/training.py", line 1646, in fit
    self._make_train_function()
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/training.py", line 970, in _make_train_function
    loss=self.total_loss)
  File "/usr/local/lib/python2.7/dist-packages/keras/legacy/interfaces.py", line 91, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/keras/optimizers.py", line 162, in get_updates
    grads = self.get_gradients(loss, params)
  File "/usr/local/lib/python2.7/dist-packages/keras/optimizers.py", line 78, in get_gradients
    grads = K.gradients(loss, params)
  File "/usr/local/lib/python2.7/dist-packages/keras/backend/tensorflow_backend.py", line 2512, in gradients
    return tf.gradients(loss, variables, colocate_gradients_with_ops=True)
  File "/home/daniel/.local/lib/python2.7/site-packages/tensorflow/python/ops/gradients_impl.py", line 609, in gradients
    grad_scope, op, func_call, lambda: grad_fn(op, *out_grads))
  File "/home/daniel/.local/lib/python2.7/site-packages/tensorflow/python/ops/gradients_impl.py", line 375, in _MaybeCompile
    return grad_fn()  # Exit early
  File "/home/daniel/.local/lib/python2.7/site-packages/tensorflow/python/ops/gradients_impl.py", line 609, in <lambda>
    grad_scope, op, func_call, lambda: grad_fn(op, *out_grads))
  File "/home/daniel/.local/lib/python2.7/site-packages/tensorflow/python/ops/array_grad.py", line 734, in _ExtractImagePatchesGrad
    cols_out = int(ceil(cols_in / stride_h))
TypeError: unsupported operand type(s) for /: 'NoneType' and 'long'

Plus, I'm not sure if what I am returning is correct.

Any help would be appreciated.

This is what I have so far:

function to calculate c*s values for SSIM:

adapted from : https://gist.github.com/Dref360/a48feaecfdb9e0609c6a02590fd1f91b

def SSIM_cs(y_true, y_pred):
    patches_true = tf.extract_image_patches(y_true, [1, 8, 8, 1], [1, 2, 2, 1], [1, 1, 1, 1], "SAME")
    patches_pred = tf.extract_image_patches(y_pred, [1, 8, 8, 1], [1, 2, 2, 1], [1, 1, 1, 1], "SAME")

    var_true = K.var(patches_true, axis=3)
    var_pred = K.var(patches_pred, axis=3)
    std_true = K.sqrt(var_true)
    std_pred = K.sqrt(var_pred)
    c2 = 0.03 ** 2
    ssim = (2 * std_pred * std_true + c2)
    denom = (var_pred + var_true + c2)
    ssim /= denom
    ssim = tf.where(tf.is_nan(ssim), K.zeros_like(ssim), ssim)
    return K.mean(ssim)

functions to obtain Gaussian Kernel

adapted from: https://github.com/keras-team/keras/issues/3720

def gaussian(x, mu, sigma):
    return np.exp(-(float(x) - float(mu)) ** 2 / (2 * sigma ** 2))


def make_kernel(sigma):
    # kernel radius = 2*sigma, but minimum 3x3 matrix
    kernel_size = max(3, int(2 * 2 * sigma + 1))
    mean = np.floor(0.5 * kernel_size)
    kernel_1d = np.array([gaussian(x, mean, sigma) for x in range(kernel_size)])
    # make 2D kernel
    np_kernel = np.outer(kernel_1d, kernel_1d).astype(dtype=K.floatx())
    # normalize kernel by sum of elements
    kernel = np_kernel / np.sum(np_kernel)
    kernel = np.reshape(kernel, (kernel_size, kernel_size, 1,1))    #height, width, in_channels, out_channel
    return kernel

Main Loss Function

def custom_Loss(y_true, y_pred):

    i   iterations = 5
weight = [0.0448, 0.2856, 0.3001, 0.2363, 0.1333]
ms_ssim = []

img1=y_true
img2=y_pred

test = []
gaussian = make_kernel(1.5)

for iteration in range(iterations):
    #Obatain c*s for current iteration
    ms_ssim.append(SSIM_cs(img1, img2)**weight[iteration])

    #Blur and Shrink
    #Transpose due to data being in order: batch, channel, height, width
    #cs for all 5 iterations -> shrink 4 times (the last is required for calculation of l)
    if(iteration!=4):
        img1 = tf.nn.conv2d(tf.transpose(img1, [0, 2, 3, 1]), gaussian, strides=[1, 1, 1, 1], padding='SAME')
        img1 = tf.transpose(img1, [0, 3, 1, 2])

        img2 = tf.nn.conv2d(tf.transpose(img2, [0, 2, 3, 1]), gaussian, strides=[1, 1, 1, 1], padding='SAME')
        img2 = tf.transpose(img2, [0, 3, 1, 2])

        img1 = K.resize_images(img1, 2,2, 'channels_first')
        img2 = K.resize_images(img2, 2,2, 'channels_first')

    ms_ssim = tf.stack(ms_ssim)
    cs_val = tf.reduce_prod(ms_ssim,0)

    patches_true = tf.extract_image_patches(img1, [1, 8, 8, 1], [1, 2, 2, 1], [1, 1, 1, 1], "SAME")
    patches_pred = tf.extract_image_patches(img2, [1, 8, 8, 1], [1, 2, 2, 1], [1, 1, 1, 1], "SAME")

    u_true = K.mean(patches_true, axis=3)
    u_pred = K.mean(patches_pred, axis=3)
    c1 = 0.01 ** 2
    l_num = (2 * u_true * u_pred + c1)
    l_den = (u_true ** 2 + u_pred ** 2 + c1)
    l_val = l_num/l_den
l_val = tf.where(tf.is_nan(l_val), K.zeros_like(l_val), l_val)
final_l_val = K.mean(l_val)


return tf.multiply(cs_val, final_l_val)

Answer 1

The problem seems to lie within tf.extract_image_patches, since this function does not allow backpropagation. You should probably create your own patch extractor using Keras backend.

TypeError: unsupported operand type(s) for /: 'NoneType' and 'long'

I noticed that this error looks like to be binded to the 'SAME' parameter, probabily triggered when calculating how many patches the function should create. In my case, changing 'SAME' to 'VALID' generated:

TypeError: unsupported operand type(s) for -: 'NoneType' and 'long'

Since the paper states:

Instead of using an 8 × 8 square window as in [3], a smooth windowing approach is used for local statistics to avoid “blocking artifacts” in the quality map [5]. Finally, a mean SSIM index of the quality map is used to evaluate the overall image quality.

I've decided to apply a convolution with a gaussian kernel and then calculate C, S and L on the resulting maps. So, at the end, my Ms_SSIM function looks like:

def keras_SSIM_cs(y_true, y_pred):
    axis=None
    gaussian = make_kernel(1.5)
    x = tf.nn.conv2d(y_true, gaussian, strides=[1, 1, 1, 1], padding='SAME')
    y = tf.nn.conv2d(y_pred, gaussian, strides=[1, 1, 1, 1], padding='SAME')

    u_x=K.mean(x, axis=axis)
    u_y=K.mean(y, axis=axis)

    var_x=K.var(x, axis=axis)
    var_y=K.var(y, axis=axis)

    cov_xy=cov_keras(x, y, axis)

    K1=0.01
    K2=0.03
    L=1  # depth of image (255 in case the image has a differnt scale)

    C1=(K1*L)**2
    C2=(K2*L)**2
    C3=C2/2

    l = ((2*u_x*u_y)+C1) / (K.pow(u_x,2) + K.pow(u_x,2) + C1)
    c = ((2*K.sqrt(var_x)*K.sqrt(var_y))+C2) / (var_x + var_y + C2)
    s = (cov_xy+C3) / (K.sqrt(var_x)*K.sqrt(var_y) + C3)

    return [c,s,l]

def keras_MS_SSIM(y_true, y_pred):
    iterations = 5
    x=y_true
    y=y_pred
    weight = [0.0448, 0.2856, 0.3001, 0.2363, 0.1333]
    c=[]
    s=[]
    for i in range(iterations):
        cs=keras_SSIM_cs(x, y)
        c.append(cs[0])
        s.append(cs[1])
        l=cs[2]
        if(i!=4):
            x=tf.image.resize_images(x, (x.get_shape().as_list()[1]//(2**(i+1)), x.get_shape().as_list()[2]//(2**(i+1))))
            y=tf.image.resize_images(y, (y.get_shape().as_list()[1]//(2**(i+1)), y.get_shape().as_list()[2]//(2**(i+1))))
    c = tf.stack(c)
    s = tf.stack(s)
    cs = c*s

    #Normalize: suggestion from https://github.com/jorge-pessoa/pytorch-msssim/issues/2 last comment to avoid NaN values
    l=(l+1)/2
    cs=(cs+1)/2

    cs=cs**weight
    cs = tf.reduce_prod(cs)
    l=l**weight[-1]

    ms_ssim = l*cs
    ms_ssim = tf.where(tf.is_nan(ms_ssim), K.zeros_like(ms_ssim), ms_ssim)

    return K.mean(ms_ssim)