Как отлаживать пользовательскую функцию потерь в Keras?

То, что я иногда делаю (точно не лучшее решение и не всегда возможно), просто заменяю бэкэнд K на np и тестирую его с некоторыми простыми данными. Вот пример

оригинальная функция Keras:

def loss(y_true, y_pred):
    means = K.reshape(y_pred[:, 0], (-1, 1))
    stds = K.reshape(y_pred[:, 1], (-1, 1))
    var = K.square(stds)
    denom = K.sqrt(2 * np.pi * var)
    prob_num = - K.square(y_true - means) / (2 * var)
    prob = prob_num - denom
    r = K.exp(prob - old_prediction)
    return -K.mean(K.minimum(r * advantage, K.clip(r, min_value=1 - self.LOSS_CLIPPING, max_value=1 + self.LOSS_CLIPPING) * advantage))

функция тестирования:

def loss(y_true, y_pred):
    means = np.reshape(y_pred[:, 0], (-1, 1))
    stds = np.reshape(y_pred[:, 1], (-1, 1))
    var = np.square(stds)
    print(var.shape)
    denom = np.sqrt(2 * np.pi * var)
    print(denom.shape)
    prob_num = - np.square(y_true - means) / (2 * var)
    prob = prob_num - denom
    r = np.exp(prob - old_prediction)
    print(r.shape)
    cliped = np.minimum(r * advantage, np.clip(r, a_min=1 - LOSS_CLIPPING, a_max=1 + LOSS_CLIPPING) * advantage)
    print(cliped.shape)
    return -np.mean(cliped)

Тестирование:

LOSS_CLIPPING = 0.2
y_pred = np.array([[2,1], [1, 1], [5, 1]])
y_true = np.array([[1], [3], [2]])
old_prediction = np.array([[-2], [-5], [-6]])
advantage = np.array([[ 0.51467506],[-0.64960159],[-0.53304715]])
loss(y_true, y_pred)

После выполнения выше, результаты:

(3, 1)
(3, 1)
(3, 1)
(3, 1)
0.43409555193679816