Как масштабировать градиентную норму в Керасе

Question

В псевдокоде для MuZero они делают следующее:

hidden_state = tf.scale_gradient(hidden_state, 0.5)

С этот вопрос о том, что это значит, я узнал, что это, вероятно, масштабирование градиентной нормы.

Как мне выполнить масштабирование нормы градиента (обрезать градиентную норму до определенной длины) в скрытом состоянии в Keras? Позже они также делают то же самое масштабирование для значения потерь:

loss += tf.scale_gradient(l, gradient_scale)

Этот сайт говорит, что я должен использовать параметр clipnorm на оптимизаторе. Но я не думаю, что это сработает, потому что я масштабирую градиенты перед использованием оптимизатора. (И тем более, что я масштабирую разные вещи на разную длину.)

Вот конкретный рассматриваемый код из статьи, на случай, если он будет полезен. (Обратите внимание, что scale_gradient не является действительной функцией Tensorflow. См. Ранее связанный вопрос , если вы не уверены, как я.)

def update_weights(optimizer: tf.train.Optimizer, network: Network, batch,
                   weight_decay: float):
  loss = 0
  for image, actions, targets in batch:
    # Initial step, from the real observation.
    value, reward, policy_logits, hidden_state = network.initial_inference(
        image)
    predictions = [(1.0, value, reward, policy_logits)]

    # Recurrent steps, from action and previous hidden state.
    for action in actions:
      value, reward, policy_logits, hidden_state = network.recurrent_inference(
          hidden_state, action)
      predictions.append((1.0 / len(actions), value, reward, policy_logits))

      # THIS LINE HERE
      hidden_state = tf.scale_gradient(hidden_state, 0.5)

    for prediction, target in zip(predictions, targets):
      gradient_scale, value, reward, policy_logits = prediction
      target_value, target_reward, target_policy = target

      l = (
          scalar_loss(value, target_value) +
          scalar_loss(reward, target_reward) +
          tf.nn.softmax_cross_entropy_with_logits(
              logits=policy_logits, labels=target_policy))

      # AND AGAIN HERE
      loss += tf.scale_gradient(l, gradient_scale)

  for weights in network.get_weights():
    loss += weight_decay * tf.nn.l2_loss(weights)

  optimizer.minimize(loss)

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

Answer 1

Вы можете использовать представленное ограничение MaxNorm здесь .

Это очень просто и понятно. Импортируйте его from keras.constraints import MaxNorm

Если вы хотите применить его к весам, когда вы определяете слой Keras, вы используете kernel_constraint = MaxNorm(max_value=2, axis=0) (подробности по оси см. На странице)

Вы также можете используйте bias_constraint = ...

Если вы хотите применить его к любому другому тензору, вы можете просто вызвать его с помощью тензора:

normalizer = MaxNorm(max_value=2, axis=0)
normalized_tensor = normalizer(original_tensor)

И вы можете увидеть исходный код довольно просто:

def __call__(self, w):
    norms = K.sqrt(K.sum(K.square(w), axis=self.axis, keepdims=True))
    desired = K.clip(norms, 0, self.max_value)
    return w * (desired / (K.epsilon() + norms))