mindscience.common.AdaHessian

class mindscience.common.AdaHessian[源代码]

二阶优化器 AdaHessian，利用 Hessian 矩阵对角元信息进行二阶优化求解。 有关更多详细信息，请参考论文 ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning 。 此处 Hessian power 固定为 1，且对 Hessian 对角元做空间平均的方法如下：

  • 对于 1D 张量：不做空间平均。

  • 对于 2D 张量：做行平均。

  • 对于 3D 张量（假设为 1D 卷积）：对最后一个维度做平均。

  • 对于 4D 张量（假设为 2D 卷积）：对最后两个维度做平均。

参数说明详见 mindspore.nn.Adam 。

样例：

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore import ops, nn
>>> from mindscience.common import AdaHessian
>>> ms.set_context(device_target="Ascend", mode=ms.GRAPH_MODE)
>>> net = nn.Conv2d(in_channels=2, out_channels=4, kernel_size=3)
>>> def forward(a):
>>>     return ops.mean(net(a)**2)**.5
>>> grad_fn = ms.grad(forward, grad_position=None, weights=net.trainable_params())
>>> optimizer = AdaHessian(net.trainable_params())
>>> inputs = ms.Tensor(np.reshape(range(100), [2, 2, 5, 5]), dtype=ms.float32)
>>> optimizer(grad_fn, inputs)
>>> print(optimizer.moment2[0].shape)
(4, 2, 3, 3)