forked from mindspore-Ecosystem/mindspore
145 lines
5.4 KiB
Python
145 lines
5.4 KiB
Python
# Copyright 2020 Huawei Technologies Co., Ltd
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ============================================================================
|
|
|
|
import pytest
|
|
import numpy as np
|
|
import mindspore.context as context
|
|
from mindspore import Tensor, Parameter
|
|
from mindspore.nn import Cell
|
|
from mindspore.nn.graph_kernels import LambUpdateWithLR, LambNextMV
|
|
|
|
|
|
class LambNet(Cell):
|
|
def __init__(self, i2, i5, x6):
|
|
super(LambNet, self).__init__()
|
|
self.i2 = Parameter(i2, name='i2')
|
|
self.i5 = Parameter(i5, name='i5')
|
|
self.x6 = Parameter(x6, name='x6')
|
|
self.lamb_next = LambNextMV()
|
|
self.lamb_update = LambUpdateWithLR()
|
|
|
|
def construct(self, i1, i3, i4, i6, i7, i8, i9, ix0, ix1, ix2, ix3,
|
|
x1, x2, x3, x4, x5, gy, se, my):
|
|
i1_ = i1 + i3
|
|
return self.lamb_next(i1_, self.i2, i3, i4, self.i5, i6, i7, i8, i9, ix0,
|
|
ix1, ix2, ix3), \
|
|
self.lamb_update(x1, x2, x3, x4, x5, self.x6, gy, se, my)
|
|
|
|
|
|
def LambUpdateNumpy(x1, x2, x3, x4, x5, x6, gy, se, my):
|
|
trust_ratio = np.where(np.greater(x2, gy),
|
|
np.where(np.greater(x1, gy), np.divide(x2, x3), se),
|
|
se)
|
|
trust_ratio = np.maximum(np.minimum(trust_ratio, my), gy)
|
|
update_with_lr = trust_ratio * x4 * x5
|
|
next_param = x6 - np.reshape(update_with_lr, x6.shape)
|
|
return next_param
|
|
|
|
|
|
def LambNextMVNumpy(i1, i2, i3, i4, i5, i6, i7, i8, i9, x0, x1, x2, x3):
|
|
m_fp32 = i5.astype(np.float32)
|
|
v_fp32 = i2.astype(np.float32)
|
|
next_m = i8 * m_fp32 + i9 * i4
|
|
next_v = x0 * v_fp32 + x1 * i1
|
|
next_mm = next_m / i6
|
|
next_vv = next_v / i3
|
|
update = next_mm / (np.sqrt(next_vv) + x3)
|
|
add3 = next_mm / np.sqrt(next_vv + x3) + x2 * i7
|
|
return add3, next_m, next_v, update
|
|
|
|
|
|
def tensor_all(*args):
|
|
res = [Tensor(a) for a in args]
|
|
return res
|
|
|
|
|
|
def test_graph_kernel_lamb():
|
|
shape = [1, 16]
|
|
oshape = [1]
|
|
np.random.seed(0)
|
|
x1 = np.random.normal(0, 1, oshape).astype(np.float32)
|
|
x2 = np.random.normal(0, 1, oshape).astype(np.float32)
|
|
x3 = np.random.normal(0, 1, oshape).astype(np.float32)
|
|
x4 = np.random.normal(0, 1, oshape).astype(np.float32)
|
|
x5 = np.random.normal(0, 1, shape).astype(np.float32)
|
|
x6 = np.random.normal(0, 1, shape).astype(np.float32)
|
|
gy = np.random.normal(0, 1, oshape).astype(np.float32)
|
|
se = np.random.normal(0, 1, oshape).astype(np.float32)
|
|
my = np.random.normal(0, 1, oshape).astype(np.float32)
|
|
|
|
tx1, tx2, tx3, tx4, tx5, tx6, tgy, tse, tmy = tensor_all(
|
|
x1, x2, x3, x4, x5, x6, gy, se, my)
|
|
|
|
np.random.seed(1)
|
|
i1 = np.abs(np.random.normal(0, 1, shape)).astype(np.float32)
|
|
i2 = np.abs(np.random.normal(0, 1, shape)).astype(np.float32)
|
|
i3 = np.abs(np.random.normal(0, 1, shape)).astype(np.float32)
|
|
i4 = np.random.normal(0, 1, shape).astype(np.float32)
|
|
i5 = np.random.normal(0, 1, shape).astype(np.float32)
|
|
i6 = np.abs(np.random.normal(0, 1, shape)).astype(np.float32)
|
|
i7 = np.random.normal(0, 1, shape).astype(np.float32)
|
|
i8 = np.random.normal(0, 1, shape).astype(np.float32)
|
|
i9 = np.random.normal(0, 1, shape).astype(np.float32)
|
|
ix0 = np.abs(np.random.normal(0, 1, shape)).astype(np.float32)
|
|
ix1 = np.abs(np.random.normal(0, 1, shape)).astype(np.float32)
|
|
ix2 = np.random.normal(0, 1, shape).astype(np.float32)
|
|
ix3 = np.ones(shape).astype(np.float32) * 1e-6
|
|
|
|
ti1, ti2, ti3, ti4, ti5, ti6, ti7, ti8, ti9, tix0, tix1, tix2, tix3 = \
|
|
tensor_all(i1, i2, i3, i4, i5, i6, i7, i8, i9, ix0, ix1, ix2, ix3)
|
|
|
|
context.set_context(enable_graph_kernel=True)
|
|
|
|
net = LambNet(ti2, ti5, tx6)
|
|
(wa3, wup), _ = net(ti1, ti3, ti4, ti6, ti7, ti8, ti9, tix0, tix1, tix2, tix3,
|
|
tx1, tx2, tx3, tx4, tx5, tgy, tse, tmy)
|
|
|
|
wi2 = net.i2.data.asnumpy().copy()
|
|
wi5 = net.i5.data.asnumpy().copy()
|
|
ares = net.x6.data.asnumpy().copy()
|
|
|
|
context.set_context(enable_graph_kernel=False)
|
|
|
|
i1_ = i1 + i3
|
|
a3, a0, a1, up = LambNextMVNumpy(i1_, i2, i3, i4, i5, i6, i7, i8, i9, ix0,
|
|
ix1, ix2, ix3)
|
|
|
|
np_res = LambUpdateNumpy(x1, x2, x3, x4, x5, x6, gy, se, my)
|
|
|
|
rtol = 0.0001
|
|
atol = 0.0001
|
|
|
|
wres = (wa3.asnumpy().copy(), wi5, wi2, wup.asnumpy().copy())
|
|
bres = (a3, a0, a1, up)
|
|
|
|
cmp_res = list(map(lambda x, y: np.allclose(x, y, rtol, atol),
|
|
wres, bres))
|
|
|
|
assert all(cmp_res) and np.allclose(ares, np_res, rtol, atol)
|
|
|
|
|
|
def test_graph_kernel_lamb_gpu():
|
|
context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="GPU")
|
|
test_graph_kernel_lamb()
|
|
|
|
|
|
@pytest.mark.level0
|
|
@pytest.mark.platform_arm_ascend_training
|
|
@pytest.mark.platform_x86_ascend_training
|
|
@pytest.mark.env_onecard
|
|
def test_graph_kernel_lamb_ascend():
|
|
context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="Ascend")
|
|
test_graph_kernel_lamb()
|