Custom loss function in tensorflow involving convolution-CodePudding

I'm trying to implement a custom loss function using convolution of an image with a kernel, very similar to what this question is doing. I have prepared my data to be in the format of (batch_size, height, width, channels). Specifically for my case, this would be (5, 500, 500, 3).

Example of a 2D kernel that I'm using:
[0 0 0; -1 0 1; 0 0 0] for channel 0
[0 1 0; 0 0 0; 0 -1 0] for channel 1
[0 0 0; 0 1 0; 0 0 0] for channel 2

I want to convolve each channel with a different kernel, and sum up the results. This is to be done for each of the 5 images in the batch. The result should be 5 500x500 image with 1 channel only, so the output that can either have shape of (batch_size, height, width, 1) or (batch_size, height, width).

To give a complete picture of what I'm trying to do, I intend to repeat the above process 2 more times with different sets of filters. So now I will have 3 outputs of shape (batch_size, height, width, 1), which I can stack to for a tensor of the shape (batch_size, height, width, 3). This is the same shape as the original tensor. After that, I am going to do an element-wise product with another tensor and sum over to calculate the loss value.

I have found a few 2D convolution functions: tf.nn.conv2d and tf.keras.backend.conv2d and tf.keras.layers.Conv2D. Which one is appropriate for this purpose? Also is it better to do the convolution using a single 3x3x3 kernel? Or perhaps 3D convolution with 3x3x3 kernel?

If you can provide a simple example or link, that would be great! Here's a template

import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras import backend as K

def my_loss(y_true, y_pred):
    kernelx0 = tf.convert_to_tensor(np.array([[0, 0, 0], [-1, 0, 1], [0, 0, 0]]))
    kernely0 = tf.convert_to_tensor(np.array([[0, 1, 0], [0, 0, 0], [0, -1, 0]]))
    kernelz0 = tf.convert_to_tensor(np.array([[0, 0, 0], [0, 1, 0], [0, 0, 0]]))

    kernelx1 = ...
    kernely1 = ...
    kernelz1 = ...

    kernelx2 = ...
    kernely2 = ...
    kernelz2 = ...

    # how to do convolution?

    return loss

CodePudding user response：

Try this way :

#this is your custom filter
kernel_in = tf.constant([[[0, 0, 0], [-1, 0, 1], [0, 0, 0]],
                  [[0, 1, 0], [0, 0, 0], [0, -1, 0]],
                  [[0, 0, 0], [0, 1, 0], [0, 0, 0]]], 
                  dtype=tf.float32)

kernel_in = tf.expand_dims(kernel_in,axis=-1) #expand this into a 4D tensor
#print(kernel_in.shape) #shape = (3,3,3,1)

out = tf.nn.conv2d(x_in, kernel_in, strides=[1, 1], padding='VALID') #x_in is your input