mindflow.core.AdaHessian

class mindflow.core.AdaHessian(params, learning_rate=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, use_locking=False, use_nesterov=False, weight_decay=0.0, loss_scale=1.0, use_amsgrad=False, **kwargs)

二阶优化器 AdaHessian，利用 Hessian 矩阵对角元信息进行二阶优化求解。 有关更多详细信息，请参考论文 ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning 。 相关 Torch 版本实现可参考 Torch 版代码 。 此处 Hessian power 固定为 1，且对 Hessian 对角元做空间平均的方法与 Torch 实现的默认行为一致，描述如下：

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

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

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

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

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

支持平台：

Ascend

样例：

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore import ops, nn
>>> from mindflow 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)