Search code examples
pythontensorflowkerasdeep-learningconvolution

Applying Gaussian blur on tensor in custom loss


I have a custom loss where I want to apply Gaussian filter to a predicted label to manipulate it a little. Using max or average pooling is simple as it is predefined in keras, but I had to make my own class for Gaussian pooling:

import numpy as np
from keras.layers import DepthwiseConv2D
from keras.layers import Input
from keras.models import Model
import tensorflow as tf

class Gaussian():
    
    def __init__(self,shape, f = 3):
        self.filt = f
        self.g = self.gaussFilter(shape)
        
    def doFilter(self, data):
        return self.g.predict(data, steps=1) #steps are for predicting on const tensor, I change it when predicting on predictions 
    
    def gauss2D(self,shape=(3,3),sigma=0.5):
    
        m,n = [(ss-1.)/2. for ss in shape]
        y,x = np.ogrid[-m:m+1,-n:n+1]
        h = np.exp( -(x*x + y*y) / (2.*sigma*sigma) )
        h[ h < np.finfo(h.dtype).eps*h.max() ] = 0
        sumh = h.sum()
        if sumh != 0:
            h /= sumh
        return h
    
    def gaussFilter(self, size=256):
        kernel_weights = self.gauss2D(shape=(self.filt,self.filt))
        
        
        in_channels = 1  # the number of input channels
        kernel_weights = np.expand_dims(kernel_weights, axis=-1)
        kernel_weights = np.repeat(kernel_weights, in_channels, axis=-1) # apply the same filter on all the input channels
        kernel_weights = np.expand_dims(kernel_weights, axis=-1)  # for shape compatibility reasons
        
        
        inp = Input(shape=(size,size,1))
        g_layer = DepthwiseConv2D(self.filt, use_bias=False, padding='same')(inp)
        model_network = Model(input=inp, output=g_layer)
        print(model_network.summary())
        model_network.layers[1].set_weights([kernel_weights])
        model_network.trainable= False
            
        return model_network

This works as expected when feeding a constant tensor to the doFilter function, an example of simple data:

a = np.array([[[1, 2, 3], [4, 5, 6], [4, 5, 6]]])
filt = Gaussian(3)
print(filt.doFilter(tf.constant(a.reshape(1,3,3,1))))

However, if I try to use this in a custom loss :

def custom_loss_no_true(input_tensor, length):
    def loss(y_true, y_pred):
        gaus_pooler = Gaussian(256, length//8)
        a = gaus_pooler.doFilter(y_pred)
        ...more stuff comes after

I get an error:

ValueError: When feeding symbolic tensors to a model, we expect the tensors to have a static batch size. Got tensor with shape: (None, 256, 256, 1)

This is as I have found caused by the fact, that I am feeding a tensor that is an output of other model, a symbolic data, not actual values (source). Thus I need to change the logic of my approach, because evaluating the tensor to feed my class would break the graph and lead to no gradient propagation within the loss (or am I incorrect?). How can I apply such convolution operation on a tensor that is an output of other model? Is it even possible? Or maybe there is a way to use it without adding the layer to the model, such as MaxPooling?


Solution

  • You don't really need a complex keras Model nor a keras Layer if what you want to do is just convolve your input with a Gaussian kernel. Here is a port of your code with simple tensorflow ops :

    import tensorflow as tf
    def get_gaussian_kernel(shape=(3,3), sigma=0.5):
        """build the gaussain filter"""
        m,n = [(ss-1.)/2. for ss in shape]
        x = tf.expand_dims(tf.range(-n,n+1,dtype=tf.float32),1)
        y = tf.expand_dims(tf.range(-m,m+1,dtype=tf.float32),0)
        h = tf.exp(tf.math.divide_no_nan(-((x*x) + (y*y)), 2*sigma*sigma))
        h = tf.math.divide_no_nan(h,tf.reduce_sum(h))
        return h
    
    def gaussian_blur(inp, shape=(3,3), sigma=0.5):
        """Convolve using tf.nn.depthwise_conv2d"""
        in_channel = tf.shape(inp)[-1]
        k = get_gaussian_kernel(shape,sigma)
        k = tf.expand_dims(k,axis=-1)
        k = tf.repeat(k,in_channel,axis=-1)
        k = tf.reshape(k, (*shape, in_channel, 1))
        # using padding same to preserve size (H,W) of the input
        conv = tf.nn.depthwise_conv2d(inp, k, strides=[1,1,1,1],padding="SAME")
        return conv
    

    You can use it simply in your custom loss (assuming a 4D y_pred [batch, height width, channel]) :

    a = gaussian_blur(y_pred)