Как вызвать функцию python внутри пользовательской функции потерь tf.keras?

Question

Внутри моей пользовательской функции потерь мне нужно вызвать функцию чистого Python, передающую вычисленные ошибки TD и некоторые индексы.Функция не должна ничего возвращать или быть дифференцированной.Вот функция, которую я хочу вызвать:

def update_priorities(self, traces_idxs, td_errors):
    """Updates the priorities of the traces with specified indexes."""
    self.priorities[traces_idxs] = td_errors + eps

Я пытался использовать tf.py_function для вызова функции-оболочки, но она вызывается только в том случае, если она встроена в график, т. Е. Если она имеет входы и выходы ивыходы используются.Поэтому я попытался пройти через некоторые из тензоров, не выполняя над ними никаких операций, и теперь вызывается функция.Вот моя полная функция потери:

def masked_q_loss(data, y_pred):
    """Computes the MSE between the Q-values of the actions that were taken and the cumulative
    discounted rewards obtained after taking those actions. Updates trace priorities.
    """
    action_batch, target_qvals, traces_idxs = data[:,0], data[:,1], data[:,2]
    seq = tf.cast(tf.range(0, tf.shape(action_batch)[0]), tf.int32)
    action_idxs = tf.transpose(tf.stack([seq, tf.cast(action_batch, tf.int32)]))
    qvals = tf.gather_nd(y_pred, action_idxs)

    def update_priorities(_qvals, _target_qvals, _traces_idxs):
        """Computes the TD error and updates memory priorities."""
        td_error = _target_qvals - _qvals
        _traces_idxs = tf.cast(_traces_idxs, tf.int32)
        mem.update_priorities(_traces_idxs, td_error)
        return _qvals

    qvals = tf.py_function(func=update_priorities, inp=[qvals, target_qvals, traces_idxs], Tout=[tf.float32])
    return tf.keras.losses.mse(qvals, target_qvals)

Однако я получаю следующую ошибку из-за вызова mem.update_priorities(_traces_idxs, td_error)

ValueError: An operation has `None` for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval.

Мне не нужно вычислять градиенты для update_prioritiesЯ просто хочу вызвать его в определенной точке вычисления графа и забыть об этом.Как я могу это сделать?

Answer 1

Использование .numpy() на тензорах внутри функции-оболочки решило проблему:

def update_priorities(_qvals, _target_qvals, _traces_idxs):
    """Computes the TD error and updates memory priorities."""
    td_error = np.abs((_target_qvals - _qvals).numpy())
    _traces_idxs = (tf.cast(_traces_idxs, tf.int32)).numpy()
    mem.update_priorities(_traces_idxs, td_error)
    return _qvals