Permutation Invariant Loss in Tensorflow - tensorflow

Question

I am working on a permutation invariant loss in Tensorflow 2.8.
The Loss takes two vectorized matrices of shape (N x 5), reshapes them (N,5) and then calculates all possible permutations (N!).
Then for all permutations a loss is calculated, and the minimum of this loss is used (for the best match).
However, I get the error message:
File "C:\Users\meist\anaconda3\envs\tf-2-8\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
raise e.with_traceback(filtered_tb) from None
File "C:\Users\meist\AppData\Roaming\Python\Python39\site-packages\tensorflow\python\framework\func_graph.py", line 1147, in autograph_handler
raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:
File "C:\Users\meist\anaconda3\envs\tf-2-8\lib\site-packages\keras\engine\training.py", line 1021, in train_function *
return step_function(self, iterator)
File "C:\Users\meist\anaconda3\envs\tf-2-8\lib\site-packages\keras\engine\training.py", line 1010, in step_function **
outputs = model.distribute_strategy.run(run_step, args=(data,))
File "C:\Users\meist\anaconda3\envs\tf-2-8\lib\site-packages\keras\engine\training.py", line 1000, in run_step **
outputs = model.train_step(data)
File "C:\Users\meist\anaconda3\envs\tf-2-8\lib\site-packages\keras\engine\training.py", line 863, in train_step
self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
File "C:\Users\meist\anaconda3\envs\tf-2-8\lib\site-packages\keras\optimizer_v2\optimizer_v2.py", line 532, in minimize
return self.apply_gradients(grads_and_vars, name=name)
File "C:\Users\meist\anaconda3\envs\tf-2-8\lib\site-packages\keras\optimizer_v2\optimizer_v2.py", line 633, in apply_gradients
grads_and_vars = optimizer_utils.filter_empty_gradients(grads_and_vars)
File "C:\Users\meist\anaconda3\envs\tf-2-8\lib\site-packages\keras\optimizer_v2\utils.py", line 73, in filter_empty_gradients
raise ValueError(f"No gradients provided for any variable: {variable}. "
ValueError: No gradients provided for any variable
Apparently there are no gradients. However, when I simply input a y_train, and y_pred I do get a loss. Here is the loss code:
from tensorflow.keras.losses import CategoricalCrossentropy
from tensorflow.keras.losses import Loss
import tensorflow as tf
from itertools import permutations
import numpy as np
import keras.backend as K
class PermInvLoss(Loss):
'''
This loss is supposed to return the minimum loss, based on the best metching of y_true and y_pred.
y_true is of dim [batchsize, Nmix x 5], and will be reshaped to [batchsize, Nmix, 5] in the call
Nmix are the number of vectors that can be permutated. The elements within the vector are fixed.
The 5 elements are [class_value,class_value,reg_value,reg_value,reg_value]
The two class values will be evaluated with CategoricalCrossentropy
The three regression values will be evaluted with MSE.
'''
def __init__(self,Nmix = 3):
super(PermInvLoss, self).__init__() # is this correct?
self.name = 'perm_inv_loss'
self.cce = CategoricalCrossentropy()
self.shape = (-1,Nmix,5) # for transforming y_true, and y_pred
variants = np.math.factorial(Nmix) # number of possible permut.
permutation_idx = list(permutations(np.arange(Nmix))) # list of permutations
perm = tf.constant(permutation_idx)
self.perm_mat = tf.constant(np.eye(Nmix)[permutation_idx],dtype = tf.float32) # permutation matrix for y_pred
eye = tf.eye(Nmix,dtype=tf.int32) # eye matrix
self.rep_mat = tf.broadcast_to(eye[tf.newaxis,...],(variants,Nmix,Nmix)) # repetition matrix for y_true
def MSE(self,y_true,y_pred,axis=(-2,-1)):
# simple MSE implementation with axis
mse = K.mean(K.square(K.abs(y_true-y_pred)),axis=axis)
return mse
def call(self, y_true, y_pred):
# reshape to [batchsize, Nmix, 5]
y_true = K.reshape(y_true,self.shape)
y_pred = K.reshape(y_pred ,self.shape)
# now y_pred is permutaed in one extra dimension (variants)
y_perm = tf.linalg.matmul(tf.cast(self.perm_mat,tf.float32),y_pred[:,tf.newaxis,...]) # [batchsize x variants x Nmix x 5]
# same for y_true, but with the repetition matrix
y_true = tf.linalg.matmul(tf.cast(self.one_mat,tf.float32),y_pred[:,tf.newaxis,...])
# print(y_perm.shape) # [batchsize x variants x Nmix x 5]
# print(y_true.shape) # [batchsize x variants x Nmix x 5]
# now we have on the second dimension all possible permutations of y_pred and can evaluate them against y_true of the same shape
# CategoricalCrossentropy for the first two values (classification)
cce = CategoricalCrossentropy(reduction='none',axis=(-1))
CE = K.sum(cce(y_true[...,:2], y_perm[...,:2]),axis=-1) # [batchsize x variants]
# MSE for other values (regression)
mse = self.MSE(y_true[...,2:], y_perm[...,2:]) # [batchsize x variants]
loss = K.min(CE+mse,axis=-1) # calculates minimum loss over the variants [batchsize]
return loss
Is the Class wrong, or is there really no Gradient?

Answer 1

I found the mistake, the line:
y_true = tf.linalg.matmul(tf.cast(self.rep_mat,tf.float32),y_pred[:,tf.newaxis,...])
should obviously be
y_true = tf.linalg.matmul(tf.cast(self.rep_mat,tf.float32),y_true[:,tf.newaxis,...])