• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15
16import numpy as np
17import pytest
18
19import mindspore.context as context
20import mindspore.nn as nn
21from mindspore.common.parameter import Parameter
22from mindspore import Tensor
23from mindspore.ops import operations as P
24
25class MomentumFusionNet(nn.Cell):
26    def __init__(self, var, accum):
27        super(MomentumFusionNet, self).__init__()
28        self.op = P.ApplyMomentum()
29        self.add = P.AddN()
30        self.mul = P.Mul()
31        self.var = Parameter(var, name="variable")
32        self.accum = Parameter(accum, name="accumulate")
33        self.lr = 0.1
34        self.weight_decay = 0.002
35        self.moment = 0.98
36
37    def construct(self, grad):
38        wd = self.mul(self.var, self.weight_decay)
39        g = self.add((wd, grad))
40        return self.op(self.var, self.accum, self.lr, g, self.moment)
41
42
43@pytest.mark.level0
44@pytest.mark.platform_x86_gpu_training
45@pytest.mark.env_onecard
46def test_momentum_fusion():
47    np.random.seed(42)
48    var = Tensor(np.random.randn(10, 20).astype(np.float32))
49    accum = Tensor(np.random.randn(10, 20).astype(np.float32))
50    grad = Tensor(np.random.randn(10, 20).astype(np.float32))
51
52    context.set_context(device_target='GPU', mode=context.GRAPH_MODE)
53    net1 = MomentumFusionNet(var, accum)
54    _ = net1(grad)
55
56    context.set_context(device_target='GPU', mode=context.PYNATIVE_MODE)
57    net2 = MomentumFusionNet(var, accum)
58    _ = net2(grad)
59
60    assert np.allclose(net1.var.data.asnumpy(), net2.var.data.asnumpy(), atol=1e-5)
61    assert np.allclose(net1.accum.data.asnumpy(), net2.accum.data.asnumpy(), atol=1e-5)
62