# Custom Loss Function in Tensorflow for UNet

I am working on a Segmentation task, where I planned to use U-Net

for the input_image of shape (224,224,3), the output should be the mask image of shape (224,224,1)

The mask image contains two unique values - black[0] and white[1]

the output layer from the UNet is having the tensor of shape (None, 224, 224, 1)

I used the softmax activation function for the output layer

The shapes and type for training data

print(trainX.shape) # (200, 224, 224, 3)
print(testX.shape) # (50, 224, 224, 3)
print(trainY.shape) # (200, 224, 224, 1)
print(testY.shape) # (50, 224, 224, 1)

print(trainX.dtype) # float64
print(testX.dtype) # float64
print(trainY.dtype) # int16
print(testY.dtype) # int16

In the mask image, the black pixel are very much more than white pixel. For balancing the black and white pixels, I planned to use class weight [for black - 0.53083749, for white - 8.60701406] in the trianing. So, I write this function

def lossFunc(true, pred):
weightsList = K.constant([0.53083749, 8.60701406])
true = K.reshape(true, [-1])
pred = K.squeeze(pred, axis=3)
sample_weightsList = K.gather(weightsList, true)
loss = keras.losses.sparse_categorical_crossentropy(true,pred)
loss*sample_weightsList

return loss

But when I started training, I got this error

InvalidArgumentError                      Traceback (most recent call last)
<timed exec> in <module>

/usr/local/lib/python3.10/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
68             # To get the full stack trace, call:
69             # tf.debugging.disable_traceback_filtering()
---> 70             raise e.with_traceback(filtered_tb) from None
71         finally:
72             del filtered_tb

/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
50   try:
51     ctx.ensure_initialized()
---> 52     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
53                                         inputs, attrs, num_outputs)
54   except core._NotOkStatusException as e:

InvalidArgumentError: No OpKernel was registered to support Op 'GatherV2' used by {{node lossFunc/GatherV2}} with these attrs: [Tparams=DT_FLOAT, Tindices=DT_INT16, batch_dims=0, Taxis=DT_INT32]
Registered devices: [CPU, GPU]
Registered kernels:
device='XLA_CPU_JIT'; Taxis in [DT_INT32, DT_INT64]; Tindices in [DT_INT32, DT_INT16, DT_INT64]; Tparams in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, 930109355527764061, DT_HALF, DT_UINT32, DT_UINT64, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN]
device='XLA_GPU_JIT'; Taxis in [DT_INT32, DT_INT64]; Tindices in [DT_INT32, DT_INT16, DT_INT64]; Tparams in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, 930109355527764061, DT_HALF, DT_UINT32, DT_UINT64, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN]
device='CPU'; Tparams in [DT_QINT16]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_QINT16]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_QUINT16]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_QUINT16]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_QINT32]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_QINT32]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_QUINT8]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_QUINT8]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_QINT8]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_QINT8]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_VARIANT]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_VARIANT]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_RESOURCE]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_RESOURCE]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_STRING]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_STRING]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_BOOL]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_BOOL]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_COMPLEX128]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_COMPLEX128]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_COMPLEX64]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_COMPLEX64]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_DOUBLE]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_DOUBLE]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_FLOAT]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_FLOAT]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_BFLOAT16]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_BFLOAT16]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_HALF]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_HALF]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_INT32]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_INT32]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_INT8]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_INT8]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_UINT8]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_UINT8]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_INT16]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_INT16]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_UINT16]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_UINT16]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_UINT32]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_UINT32]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_INT64]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_INT64]; Tindices in [DT_INT32]
device='CPU'; Tparams in [DT_UINT64]; Tindices in [DT_INT64]
device='CPU'; Tparams in [DT_UINT64]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_BOOL]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_BOOL]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_COMPLEX128]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_COMPLEX128]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_COMPLEX64]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_COMPLEX64]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_DOUBLE]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_DOUBLE]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_FLOAT]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_FLOAT]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_BFLOAT16]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_BFLOAT16]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_HALF]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_HALF]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_INT64]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_INT64]; Tindices in [DT_INT32]
device='GPU'; Tparams in [DT_INT32]; Tindices in [DT_INT64]
device='GPU'; Tparams in [DT_INT32]; Tindices in [DT_INT32]

[[lossFunc/GatherV2]] [Op:__inference_train_function_6870]

What is this issue and How to solve this

Thanks