Tensorflow custom loss function: ValueError using tf.Varaible

Question

so I am writing a custom loss function for an LSTM model in keras and tensorflow. The problem is that my loss function currently includes a tf.Variable in order to create a tensor that is repeadly updated. However, I get an error which I cannot manage to work around. Can anyone help me? Here is the code:

def custom_loss(y_true, y_pred):

    y_true_next = y_true[1:]
    y_pred_next = y_pred[1:]
    
    y_true_tdy = y_true[:-1]
    y_pred_tdy = y_pred[:-1]
    
    print('Shape of y_pred_back -', y_pred_tdy.get_shape())

    y_true_diff = tf.subtract(y_true_next, y_true_tdy)
    y_pred_diff = tf.subtract(y_pred_next, y_pred_tdy)
        
    standard = tf.zeros_like(y_pred_diff)
    
    y_true_move = tf.greater_equal(y_true_diff, standard)
    y_pred_move = tf.greater_equal(y_pred_diff, standard)
    y_true_move = tf.reshape(y_true_move, [-1])
    y_pred_move = tf.reshape(y_pred_move, [-1])
    
    condition = tf.not_equal(y_true_move, y_pred_move)
    indices = tf.where(condition)

    ones = tf.ones_like(indices)
    indices = tf.add(indices, ones)
    indices = K.cast(indices, dtype='int32')
    
    direction_loss = tf.Variable(tf.ones_like(y_pred), dtype='float32')
    updates = K.cast(tf.ones_like(indices), dtype='float32')
    alpha = 1000
    direction_loss = tf.compat.v1.scatter_nd_update(direction_loss, indices, alpha*updates)
    
    custom_loss = K.mean(tf.multiply(K.square(y_true - y_pred), direction_loss), axis=-1)
    
    return custom_loss

And the error states:

ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information.

I will include the rest of the model definition and the training call

def create_lstm_model():
      
    lstm_model = Sequential()
    lstm_model.add(LSTM(100, batch_input_shape=(BATCH_SIZE, TIME_STEPS, x_t.shape[2]),
                        dropout=0.0, recurrent_dropout=0.0,
                        stateful=True, return_sequences=True,
                        kernel_initializer='random_uniform'))

    lstm_model.add(LSTM(60, dropout=0.0))
    
    lstm_model.add(Dense(20,activation='relu'))
    lstm_model.add(Dense(1, activation='sigmoid'))
        
    #compile the model
    optimizer = tf.keras.optimizers.Adam(lr=params["LR"])
    lstm_model.compile(loss=custom_loss, optimizer=optimizer)
    print('model check')
  
    return lstm_model
    lstm_model = create_lstm_model()
    print(lstm_model.summary())
    
history_lstm = lstm_model.fit(x_t, y_t, epochs=params["EPOCHS"], verbose=1, batch_size=BATCH_SIZE,
                      shuffle=False, validation_data=(trim_dataset(x_val, BATCH_SIZE),
                      trim_dataset(y_val, BATCH_SIZE)))

user_12926704 · Answer

I stucked on that problem a while. It seems that you should not use tf variable in loss function ;therefore, I found a solution without using scatter_nd_update function which requires tf.variable. You can think scatter as direction_loss variable.

    scatter = tf.scatter_nd(tf.cast(indices, tf.int32), updates, tf.shape(y_pred)) * alpha
    scatter = tf.where(tf.equal(scatter, 0), tf.ones_like(scatter), scatter)
    scatter = K.cast(scatter, dtype='float32')

Tensorflow custom loss function: ValueError using tf.Varaible

1 Answers