python - 带状态的自定义渐变

Question

我正在尝试在 tensorflow中实现这个梯度剪报，这需要存储梯度规范的历史。

我假设我需要使用tf.custom_gradient装饰器来执行此操作，但是如何维护梯度规范历史的运行列表？我可以像在 pytorch 版本中那样使用闭包吗？

作为参考，这里是pytorch中的实现。

import numpy as np
import torch
from enum import Enum

def _get_grad_norm(model):
    total_norm = 0
    for p in model.parameters():
        if p.grad is not None:
            param_norm = p.grad.data.norm(2)
            total_norm += param_norm.item() ** 2
    total_norm = total_norm ** (1. / 2)
    return total_norm 

# written for pytorch ignite
# fire this on backwards pass
class BackwardsEvents(Enum):
    BACKWARDS_COMPLETED = 'backwards_completed'

def add_autoclip_gradient_handler(engine, model, clip_percentile):
    # Keep track of the history of gradients and select a cutoff
    # to clip values to based on percentile.
    grad_history = []

    @engine.on(BackwardsEvents.BACKWARDS_COMPLETED)
    def autoclip_gradient(engine):
        obs_grad_norm = _get_grad_norm(model)
        grad_history.append(obs_grad_norm)
        clip_value = np.percentile(grad_history, clip_percentile)
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_value)

标签： pythontensorflowpytorchgradient-descent