简化Pythorch自动加载计算的库
autograd-lib的Python项目详细描述
自动签名
作者:Yaroslav Bulatov,Kazuki Osawa
库来简化Pythorch中的梯度计算。在
示例1:每个示例的梯度规范
使用它计算线性层的梯度范数的示例,使用https://arxiv.org/abs/1510.01799的技巧
有关可运行的示例,请参见example_norms.py
。重要部分:
!pip install autograd-lib
from autograd_lib import autograd_lib
loss_fn = ...
data = ...
model = ...
autograd_lib.register(model)
activations = {}
def save_activations(layer, A, _):
activations[layer] = A
with autograd_lib.module_hook(save_activations):
output = model(data)
loss = loss_fn(output)
norms = [torch.zeros(n)]
def per_example_norms(layer, _, B):
A = activations[layer]
norms[0]+=(A*A).sum(dim=1)*(B*B).sum(dim=1)
with autograd_lib.module_hook(per_example_norms):
loss.backward()
print('per-example gradient norms squared:', norms[0])
例2:黑森量
在一次迭代中计算ReLU网络所有线性层的精确Hessian、Hessian对角线和KFAC近似。在
请参见example_hessian.py
以获取自包含的示例。重要部分:
变化:
autograd_lib.backward_hessian
适用于黑森autograd_lib.backward_jacobian
表示雅可比平方- ^经验Fisher信息矩阵的{
}
见自动签名_测试.py根据Pythorch autograd进行正确性检查。在
- 项目
标签: