Add ut for cpu kernel of sparse optimizer

mindspore/ccsrc/kernel/cpu/sparse_apply_adam_cpu_kernel.cc

@@ -81,6 +81,7 @@ void SparseApplyAdamCPUKernel::InitInputOutputSize(const CNodePtr &kernel_node)
   MS_EXCEPTION_IF_NULL(kernel_node);
   workspace_size_list_.emplace_back(indices_size_ * var_outer_dim_size_ * sizeof(float));
   workspace_size_list_.emplace_back(indices_size_ * sizeof(int));
+  workspace_size_list_.emplace_back(var_first_dim_size_ * var_outer_dim_size_ * sizeof(float));
 }
 
 void SparseApplyAdamCPUKernel::InitKernel(const CNodePtr &kernel_node) {
@@ -141,6 +142,7 @@ bool SparseApplyAdamCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inp
   auto indices = reinterpret_cast<int *>(inputs[10]->addr);
   auto new_grad = reinterpret_cast<float *>(workspace[0]->addr);
   auto new_indices = reinterpret_cast<int *>(workspace[1]->addr);
+  auto m_t = reinterpret_cast<float *>(workspace[2]->addr);
 
   SparseGradient unique_sparse_grad({new_grad, new_indices, indices_size_});
   ReduceSparseGradient(SparseGradient({grad, indices, indices_size_}), &unique_sparse_grad, var_first_dim_size_,
@@ -156,8 +158,7 @@ bool SparseApplyAdamCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inp
   const size_t kThreadNum = 16;
   MultiThreadCompute(ComputeMomentum, &input_params, kThreadNum, total_dim_size);
 
-  std::vector<float> m_t(m, m + total_dim_size);
-  input_params.m_t_ = m_t.data();
+  input_params.m_t_ = m_t;
   input_params.use_nesterov_ = use_nesterov_;
   input_params.sparse_grad_ = unique_sparse_grad;
   input_params.var_first_dim_size_ = var_first_dim_size_;

tests/st/ops/cpu/test_sparse_apply_adam_op.py

@@ -44,10 +44,20 @@ class Net(nn.Cell):
 
 
 def test_net():
-    gradient = Tensor(np.random.rand(3, 3, 3).astype(np.float32))
+    gradient = Tensor(np.ones([3, 3, 3]).astype(np.float32))
     indices = Tensor([0, 1, 2], mstype.int32)
 
     context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
     sparse_apply_adam = Net()
-    output = sparse_apply_adam(gradient, indices)
-    print(output[0].asnumpy())
+    sparse_apply_adam(gradient, indices)
+    print(sparse_apply_adam.var.default_input)
+    expect_var = np.array([[[0.9996838, 0.9996838, 0.9996838],
+                            [0.9996838, 0.9996838, 0.9996838],
+                            [0.9996838, 0.9996838, 0.9996838]],
+                           [[0.9996838, 0.9996838, 0.9996838],
+                            [0.9996838, 0.9996838, 0.9996838],
+                            [0.9996838, 0.9996838, 0.9996838]],
+                           [[0.9996838, 0.9996838, 0.9996838],
+                            [0.9996838, 0.9996838, 0.9996838],
+                            [0.9996838, 0.9996838, 0.9996838]]]).astype(np.float32)
+    assert np.all(sparse_apply_adam.var.default_input.asnumpy() == expect_var)

tests/st/ops/cpu/test_sparse_apply_ftrl_op.py

@@ -36,15 +36,20 @@ class Net(nn.Cell):
 
 
 def test_net():
-    gradient = Tensor(np.random.rand(3, 3, 3).astype(np.float32))
+    gradient = Tensor(np.ones([3, 3, 3]).astype(np.float32))
     indices = Tensor([0, 1, 2], mstype.int32)
 
     context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
     sparse_apply_ftrl = Net()
-    output = sparse_apply_ftrl(gradient, indices)
-    print(output[0].asnumpy())
-
-    context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
-    sparse_apply_ftrl = Net()
-    output = sparse_apply_ftrl(gradient, indices)
-    print(output[0].asnumpy())
+    sparse_apply_ftrl(gradient, indices)
+    print(sparse_apply_ftrl.var.default_input)
+    expect_var = np.array([[[0.291479, 0.291479, 0.291479],
+                            [0.291479, 0.291479, 0.291479],
+                            [0.291479, 0.291479, 0.291479]],
+                           [[0.291479, 0.291479, 0.291479],
+                            [0.291479, 0.291479, 0.291479],
+                            [0.291479, 0.291479, 0.291479]],
+                           [[0.291479, 0.291479, 0.291479],
+                            [0.291479, 0.291479, 0.291479],
+                            [0.291479, 0.291479, 0.291479]]]).astype(np.float32)
+    assert np.all(sparse_apply_ftrl.var.default_input.asnumpy() == expect_var)

tests/st/ops/cpu/test_sparse_apply_proximal_adagrad_op.py

@@ -38,10 +38,20 @@ class Net(nn.Cell):
 
 
 def test_net():
-    gradient = Tensor(np.random.rand(3, 3, 3).astype(np.float32))
+    gradient = Tensor(np.ones([3, 3, 3]).astype(np.float32))
     indices = Tensor([0, 1, 2], mstype.int32)
 
     context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
     sparse_apply_proximal_adagrad = Net()
-    output = sparse_apply_proximal_adagrad(gradient, indices)
-    print(output[0].asnumpy())
+    sparse_apply_proximal_adagrad(gradient, indices)
+    print(sparse_apply_proximal_adagrad.var.default_input)
+    expect_var = np.array([[[0.9929289, 0.9929289, 0.9929289],
+                            [0.9929289, 0.9929289, 0.9929289],
+                            [0.9929289, 0.9929289, 0.9929289]],
+                           [[0.9929289, 0.9929289, 0.9929289],
+                            [0.9929289, 0.9929289, 0.9929289],
+                            [0.9929289, 0.9929289, 0.9929289]],
+                           [[0.9929289, 0.9929289, 0.9929289],
+                            [0.9929289, 0.9929289, 0.9929289],
+                            [0.9929289, 0.9929289, 0.9929289]]]).astype(np.float32)
+    assert np.all(sparse_apply_proximal_adagrad.var.default_input.asnumpy() == expect_var)

tests/ut/cpp/CMakeLists.txt

@@ -104,6 +104,12 @@ file(GLOB_RECURSE MINDSPORE_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR}
         "../../../mindspore/ccsrc/predict/converter/attr_utils/*.cc"
         "../../../mindspore/ccsrc/predict/converter/lite_model/*.cc"
         "../../../mindspore/ccsrc/predict/converter/lite_model/operations/*.cc"
+        "../../../mindspore/ccsrc/kernel/cpu/cpu_kernel.cc"
+        "../../../mindspore/ccsrc/kernel/cpu/cpu_kernel_factory.cc"
+        "../../../mindspore/ccsrc/kernel/cpu/sparse_apply_adam_cpu_kernel.cc"
+        "../../../mindspore/ccsrc/kernel/cpu/sparse_apply_ftrl_cpu_kernel.cc"
+        "../../../mindspore/ccsrc/kernel/cpu/sparse_apply_lazy_adam_cpu_kernel.cc"
+        "../../../mindspore/ccsrc/kernel/cpu/sparse_apply_proximal_adagrad_cpu_kernel.cc"
         )
 
 list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/debug/dump_proto.cc")

tests/ut/cpp/kernel/common_utils_test.cc

@@ -49,7 +49,7 @@ TEST_F(CommonUtilTest, DeduplicateIndexedSlicesTest1) {
   std::vector<int> unique_indices(3);
   std::vector<float> summed_grad(6);
   SparseGradient unique_grad({summed_grad.data(), unique_indices.data(), 0});
-  DeduplicateIndexedSlices(SparseGradient({grad.data(), indices.data(), 6}), &unique_grad, 6, 2);
+  ReduceSparseGradient(SparseGradient({grad.data(), indices.data(), 6}), &unique_grad, 6, 2);
   EXPECT_EQ(unique_grad.indices_size_, 3);
   EXPECT_EQ(unique_indices, std::vector<int>({0, 1, 3}));
   /* 10 13
@@ -83,7 +83,7 @@ TEST_F(CommonUtilTest, DeduplicateIndexedSlicesTest2) {
   std::vector<int> unique_indices(2);
   std::vector<float> summed_grad(4);
   SparseGradient unique_grad({summed_grad.data(), unique_indices.data(), 0});
-  DeduplicateIndexedSlices(SparseGradient({grad.data(), indices.data(), 6}), &unique_grad, 6, 2);
+  ReduceSparseGradient(SparseGradient({grad.data(), indices.data(), 6}), &unique_grad, 6, 2);
   EXPECT_EQ(unique_grad.indices_size_, 2);
   EXPECT_EQ(unique_indices, std::vector<int>({0, 1}));
   /* 10 13

tests/ut/cpp/kernel/cpu/sparse_apply_adam_cpu_kernel_test.cc

+/**
+ * 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.
+ */
+
+#include <vector>
+#include "common/common_test.h"
+#define private public
+#define protected public
+#include "kernel/cpu/sparse_apply_adam_cpu_kernel.h"
+#undef private
+#undef protected
+
+namespace mindspore {
+namespace kernel {
+class SparseApplyAdamCpuKernelTest : public UT::Common {
+ public:
+  SparseApplyAdamCpuKernelTest() : sparse_adam_(std::make_shared<SparseApplyAdamCPUKernel>()) {}
+
+  void SetUp() override {
+    var_.clear();
+    m_.clear();
+    v_.clear();
+    grad_.clear();
+    inputs_.clear();
+    workspace_.clear();
+    outputs_.clear();
+  }
+
+  AddressPtr CreateKernelAddress(void *addr) {
+    auto kernel_addr = std::make_shared<Address>();
+    kernel_addr->addr = addr;
+    return kernel_addr;
+  }
+
+  void CreateInputAddress(std::vector<int> &indices) {
+    inputs_.push_back(CreateKernelAddress(var_.data()));
+    inputs_.push_back(CreateKernelAddress(m_.data()));
+    inputs_.push_back(CreateKernelAddress(v_.data()));
+    inputs_.push_back(CreateKernelAddress(&beta1_power_));
+    inputs_.push_back(CreateKernelAddress(&beta2_power_));
+    inputs_.push_back(CreateKernelAddress(&lr_));
+    inputs_.push_back(CreateKernelAddress(&beta1_));
+    inputs_.push_back(CreateKernelAddress(&beta2_));
+    inputs_.push_back(CreateKernelAddress(&epsilon_));
+    inputs_.push_back(CreateKernelAddress(grad_.data()));
+    inputs_.push_back(CreateKernelAddress(indices.data()));
+  }
+
+  void CreateWorkspaceAddress(std::vector<float> &new_grad, std::vector<int> &new_indices, std::vector<float> &m_t) {
+    workspace_.push_back(CreateKernelAddress(new_grad.data()));
+    workspace_.push_back(CreateKernelAddress(new_indices.data()));
+    workspace_.push_back(CreateKernelAddress(m_t.data()));
+  }
+
+  std::vector<float> var_;
+  std::vector<float> m_;
+  std::vector<float> v_;
+  std::vector<float> grad_;
+  std::vector<AddressPtr> inputs_;
+  std::vector<AddressPtr> workspace_;
+  std::vector<AddressPtr> outputs_;
+  std::shared_ptr<SparseApplyAdamCPUKernel> sparse_adam_;
+  float beta1_power_ = 0.9;
+  float beta2_power_ = 0.999;
+  float lr_ = 0.001;
+  float beta1_ = 0.9;
+  float beta2_ = 0.999;
+  float epsilon_ = 1e-8;
+};
+
+TEST_F(SparseApplyAdamCpuKernelTest, dense_test) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    m_.push_back(1.0);
+    v_.push_back(1.0);
+    grad_.push_back(1.0);
+  }
+  sparse_adam_->indices_size_ = 3;
+  sparse_adam_->var_first_dim_size_ = 3;
+  sparse_adam_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{0, 1, 2};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  std::vector<float> m_t(3 * 3 * 3);
+  CreateWorkspaceAddress(new_grad, new_indices, m_t);
+  sparse_adam_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999684) < 1e-6);
+  }
+}
+
+TEST_F(SparseApplyAdamCpuKernelTest, sparse_test1) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    m_.push_back(1.0);
+    v_.push_back(1.0);
+  }
+  for (size_t i = 0; i < 2 * 3 * 3; ++i) {
+    grad_.push_back(1.0);
+  }
+  sparse_adam_->indices_size_ = 2;
+  sparse_adam_->var_first_dim_size_ = 3;
+  sparse_adam_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{0, 2};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  std::vector<float> m_t(3 * 3 * 3);
+  CreateWorkspaceAddress(new_grad, new_indices, m_t);
+  sparse_adam_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999684) < 1e-6);
+  }
+  for (size_t i = 3 * 3; i < 2 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999715) < 1e-6);
+  }
+  for (size_t i = 2 * 3 * 3; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999684) < 1e-6);
+  }
+}
+
+TEST_F(SparseApplyAdamCpuKernelTest, sparse_test2) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    m_.push_back(1.0);
+    v_.push_back(1.0);
+    grad_.push_back(1.0);
+  }
+  sparse_adam_->indices_size_ = 3;
+  sparse_adam_->var_first_dim_size_ = 3;
+  sparse_adam_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{2, 2, 1};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  std::vector<float> m_t(3 * 3 * 3);
+  CreateWorkspaceAddress(new_grad, new_indices, m_t);
+  sparse_adam_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999715) < 1e-6);
+  }
+  for (size_t i = 3 * 3; i < 2 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999684) < 1e-6);
+  }
+  for (size_t i = 2 * 3 * 3; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999653) < 1e-6);
+  }
+}
+}  // namespace kernel
+}  // namespace mindspore

tests/ut/cpp/kernel/cpu/sparse_apply_ftrl_cpu_kernel_test.cc

+/**
+ * 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.
+ */
+
+#include <vector>
+#include "common/common_test.h"
+#define private public
+#define protected public
+#include "kernel/cpu/sparse_apply_ftrl_cpu_kernel.h"
+#undef private
+#undef protected
+
+namespace mindspore {
+namespace kernel {
+class SparseApplyFtrlCpuKernelTest : public UT::Common {
+ public:
+  SparseApplyFtrlCpuKernelTest() : sparse_ftrl_(std::make_shared<SparseApplyFtrlCPUKernel>()) {}
+
+  void SetUp() override {
+    sparse_ftrl_->lr_ = 0.001;
+    sparse_ftrl_->l1_ = 0.0;
+    sparse_ftrl_->l2_ = 0.0;
+    sparse_ftrl_->lr_power_ = -0.5;
+    var_.clear();
+    accum_.clear();
+    linear_.clear();
+    grad_.clear();
+    inputs_.clear();
+    workspace_.clear();
+    outputs_.clear();
+  }
+
+  AddressPtr CreateKernelAddress(void *addr) {
+    auto kernel_addr = std::make_shared<Address>();
+    kernel_addr->addr = addr;
+    return kernel_addr;
+  }
+
+  void CreateInputAddress(std::vector<int> &indices) {
+    inputs_.push_back(CreateKernelAddress(var_.data()));
+    inputs_.push_back(CreateKernelAddress(accum_.data()));
+    inputs_.push_back(CreateKernelAddress(linear_.data()));
+    inputs_.push_back(CreateKernelAddress(grad_.data()));
+    inputs_.push_back(CreateKernelAddress(indices.data()));
+  }
+
+  void CreateWorkspaceAddress(std::vector<float> &new_grad, std::vector<int> &new_indices) {
+    workspace_.push_back(CreateKernelAddress(new_grad.data()));
+    workspace_.push_back(CreateKernelAddress(new_indices.data()));
+  }
+
+  std::vector<float> var_;
+  std::vector<float> accum_;
+  std::vector<float> linear_;
+  std::vector<float> grad_;
+  std::vector<AddressPtr> inputs_;
+  std::vector<AddressPtr> workspace_;
+  std::vector<AddressPtr> outputs_;
+  std::shared_ptr<SparseApplyFtrlCPUKernel> sparse_ftrl_;
+};
+
+TEST_F(SparseApplyFtrlCpuKernelTest, dense_test) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    accum_.push_back(1.0);
+    linear_.push_back(1.0);
+    grad_.push_back(1.0);
+  }
+  sparse_ftrl_->indices_size_ = 3;
+  sparse_ftrl_->var_first_dim_size_ = 3;
+  sparse_ftrl_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{0, 1, 2};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_ftrl_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.291479) < 1e-6);
+  }
+}
+
+TEST_F(SparseApplyFtrlCpuKernelTest, sparse_test1) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    accum_.push_back(1.0);
+    linear_.push_back(1.0);
+  }
+  for (size_t i = 0; i < 2 * 3 * 3; ++i) {
+    grad_.push_back(1.0);
+  }
+  sparse_ftrl_->indices_size_ = 2;
+  sparse_ftrl_->var_first_dim_size_ = 3;
+  sparse_ftrl_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{0, 2};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_ftrl_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.291479) < 1e-6);
+  }
+  for (size_t i = 3 * 3; i < 2 * 3 * 3; ++i) {
+    EXPECT_EQ(var_[i], 1.0);
+  }
+  for (size_t i = 2 * 3 * 3; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.291479) < 1e-6);
+  }
+}
+
+TEST_F(SparseApplyFtrlCpuKernelTest, sparse_test2) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    accum_.push_back(1.0);
+    linear_.push_back(1.0);
+    grad_.push_back(1.0);
+  }
+  sparse_ftrl_->indices_size_ = 3;
+  sparse_ftrl_->var_first_dim_size_ = 3;
+  sparse_ftrl_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{2, 2, 1};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_ftrl_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3; ++i) {
+    EXPECT_EQ(var_[i], 1.0);
+  }
+  for (size_t i = 3 * 3; i < 2 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.291479) < 1e-6);
+  }
+  for (size_t i = 2 * 3 * 3; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.551445) < 1e-6);
+  }
+}
+}  // namespace kernel
+}  // namespace mindspore

tests/ut/cpp/kernel/cpu/sparse_apply_lazy_adam_cpu_kernel_test.cc

+/**
+ * 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.
+ */
+
+#include <vector>
+#include "common/common_test.h"
+#define private public
+#define protected public
+#include "kernel/cpu/sparse_apply_lazy_adam_cpu_kernel.h"
+#undef private
+#undef protected
+
+namespace mindspore {
+namespace kernel {
+class SparseApplyLazyAdamCpuKernelTest : public UT::Common {
+ public:
+  SparseApplyLazyAdamCpuKernelTest() : sparse_lazy_adam_(std::make_shared<SparseApplyLazyAdamCPUKernel>()) {}
+
+  void SetUp() override {
+    var_.clear();
+    m_.clear();
+    v_.clear();
+    grad_.clear();
+    inputs_.clear();
+    workspace_.clear();
+    outputs_.clear();
+  }
+
+  AddressPtr CreateKernelAddress(void *addr) {
+    auto kernel_addr = std::make_shared<Address>();
+    kernel_addr->addr = addr;
+    return kernel_addr;
+  }
+
+  void CreateInputAddress(std::vector<int> &indices) {
+    inputs_.push_back(CreateKernelAddress(var_.data()));
+    inputs_.push_back(CreateKernelAddress(m_.data()));
+    inputs_.push_back(CreateKernelAddress(v_.data()));
+    inputs_.push_back(CreateKernelAddress(&beta1_power_));
+    inputs_.push_back(CreateKernelAddress(&beta2_power_));
+    inputs_.push_back(CreateKernelAddress(&lr_));
+    inputs_.push_back(CreateKernelAddress(&beta1_));
+    inputs_.push_back(CreateKernelAddress(&beta2_));
+    inputs_.push_back(CreateKernelAddress(&epsilon_));
+    inputs_.push_back(CreateKernelAddress(grad_.data()));
+    inputs_.push_back(CreateKernelAddress(indices.data()));
+  }
+
+  void CreateWorkspaceAddress(std::vector<float> &new_grad, std::vector<int> &new_indices) {
+    workspace_.push_back(CreateKernelAddress(new_grad.data()));
+    workspace_.push_back(CreateKernelAddress(new_indices.data()));
+  }
+
+  std::vector<float> var_;
+  std::vector<float> m_;
+  std::vector<float> v_;
+  std::vector<float> grad_;
+  std::vector<AddressPtr> inputs_;
+  std::vector<AddressPtr> workspace_;
+  std::vector<AddressPtr> outputs_;
+  std::shared_ptr<SparseApplyLazyAdamCPUKernel> sparse_lazy_adam_;
+  float beta1_power_ = 0.9;
+  float beta2_power_ = 0.999;
+  float lr_ = 0.001;
+  float beta1_ = 0.9;
+  float beta2_ = 0.999;
+  float epsilon_ = 1e-8;
+};
+
+TEST_F(SparseApplyLazyAdamCpuKernelTest, dense_test) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    m_.push_back(1.0);
+    v_.push_back(1.0);
+    grad_.push_back(1.0);
+  }
+  sparse_lazy_adam_->indices_size_ = 3;
+  sparse_lazy_adam_->var_first_dim_size_ = 3;
+  sparse_lazy_adam_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{0, 1, 2};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_lazy_adam_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999684) < 1e-6);
+  }
+}
+
+TEST_F(SparseApplyLazyAdamCpuKernelTest, sparse_test1) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    m_.push_back(1.0);
+    v_.push_back(1.0);
+  }
+  for (size_t i = 0; i < 2 * 3 * 3; ++i) {
+    grad_.push_back(1.0);
+  }
+  sparse_lazy_adam_->indices_size_ = 2;
+  sparse_lazy_adam_->var_first_dim_size_ = 3;
+  sparse_lazy_adam_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{0, 2};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_lazy_adam_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999684) < 1e-6);
+  }
+  for (size_t i = 3 * 3; i < 2 * 3 * 3; ++i) {
+    EXPECT_EQ(var_[i], 1.0);
+  }
+  for (size_t i = 2 * 3 * 3; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999684) < 1e-6);
+  }
+}
+
+TEST_F(SparseApplyLazyAdamCpuKernelTest, sparse_test2) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    m_.push_back(1.0);
+    v_.push_back(1.0);
+    grad_.push_back(1.0);
+  }
+  sparse_lazy_adam_->indices_size_ = 3;
+  sparse_lazy_adam_->var_first_dim_size_ = 3;
+  sparse_lazy_adam_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{2, 2, 1};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_lazy_adam_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3; ++i) {
+    EXPECT_EQ(var_[i], 1.0);
+  }
+  for (size_t i = 3 * 3; i < 2 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999684) < 1e-6);
+  }
+  for (size_t i = 2 * 3 * 3; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.999653) < 1e-6);
+  }
+}
+}  // namespace kernel
+}  // namespace mindspore

tests/ut/cpp/kernel/cpu/sparse_apply_proximal_adagrad_cpu_kernel_test.cc

+/**
+ * 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.
+ */
+
+#include <vector>
+#include "common/common_test.h"
+#define private public
+#define protected public
+#include "kernel/cpu/sparse_apply_proximal_adagrad_cpu_kernel.h"
+#undef private
+#undef protected
+
+namespace mindspore {
+namespace kernel {
+class SparseApplyProximalAdagradCpuKernelTest : public UT::Common {
+ public:
+  SparseApplyProximalAdagradCpuKernelTest()
+      : sparse_proximal_adagrad_(std::make_shared<SparseApplyProximalAdagradCPUKernel>()) {}
+
+  void SetUp() override {
+    var_.clear();
+    accum_.clear();
+    grad_.clear();
+    inputs_.clear();
+    workspace_.clear();
+    outputs_.clear();
+  }
+
+  AddressPtr CreateKernelAddress(void *addr) {
+    auto kernel_addr = std::make_shared<Address>();
+    kernel_addr->addr = addr;
+    return kernel_addr;
+  }
+
+  void CreateInputAddress(std::vector<int> &indices) {
+    inputs_.push_back(CreateKernelAddress(var_.data()));
+    inputs_.push_back(CreateKernelAddress(accum_.data()));
+    inputs_.push_back(CreateKernelAddress(&lr_));
+    inputs_.push_back(CreateKernelAddress(&l1_));
+    inputs_.push_back(CreateKernelAddress(&l2_));
+    inputs_.push_back(CreateKernelAddress(grad_.data()));
+    inputs_.push_back(CreateKernelAddress(indices.data()));
+  }
+
+  void CreateWorkspaceAddress(std::vector<float> &new_grad, std::vector<int> &new_indices) {
+    workspace_.push_back(CreateKernelAddress(new_grad.data()));
+    workspace_.push_back(CreateKernelAddress(new_indices.data()));
+  }
+
+  std::vector<float> var_;
+  std::vector<float> accum_;
+  std::vector<float> grad_;
+  std::vector<AddressPtr> inputs_;
+  std::vector<AddressPtr> workspace_;
+  std::vector<AddressPtr> outputs_;
+  std::shared_ptr<SparseApplyProximalAdagradCPUKernel> sparse_proximal_adagrad_;
+  float lr_ = 0.01;
+  float l1_ = 0.0;
+  float l2_ = 0.0;
+};
+
+TEST_F(SparseApplyProximalAdagradCpuKernelTest, dense_test) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    accum_.push_back(1.0);
+    grad_.push_back(1.0);
+  }
+  sparse_proximal_adagrad_->indices_size_ = 3;
+  sparse_proximal_adagrad_->var_first_dim_size_ = 3;
+  sparse_proximal_adagrad_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{0, 1, 2};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_proximal_adagrad_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.9929289) < 1e-6);
+  }
+}
+
+TEST_F(SparseApplyProximalAdagradCpuKernelTest, sparse_test1) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    accum_.push_back(1.0);
+  }
+  for (size_t i = 0; i < 2 * 3 * 3; ++i) {
+    grad_.push_back(1.0);
+  }
+  sparse_proximal_adagrad_->indices_size_ = 2;
+  sparse_proximal_adagrad_->var_first_dim_size_ = 3;
+  sparse_proximal_adagrad_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{0, 2};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_proximal_adagrad_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.9929289) < 1e-6);
+  }
+  for (size_t i = 3 * 3; i < 2 * 3 * 3; ++i) {
+    EXPECT_EQ(var_[i], 1.0);
+  }
+  for (size_t i = 2 * 3 * 3; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.9929289) < 1e-6);
+  }
+}
+
+TEST_F(SparseApplyProximalAdagradCpuKernelTest, sparse_test2) {
+  for (size_t i = 0; i < 3 * 3 * 3; ++i) {
+    var_.push_back(1.0);
+    accum_.push_back(1.0);
+    grad_.push_back(1.0);
+  }
+  sparse_proximal_adagrad_->indices_size_ = 3;
+  sparse_proximal_adagrad_->var_first_dim_size_ = 3;
+  sparse_proximal_adagrad_->var_outer_dim_size_ = 9;
+
+  std::vector<int> indices{2, 2, 1};
+  CreateInputAddress(indices);
+  std::vector<float> new_grad(3 * 3 * 3);
+  std::vector<int> new_indices(3);
+  CreateWorkspaceAddress(new_grad, new_indices);
+  sparse_proximal_adagrad_->Launch(inputs_, workspace_, outputs_);
+  for (size_t i = 0; i < 3 * 3; ++i) {
+    EXPECT_EQ(var_[i], 1.0);
+  }
+  for (size_t i = 3 * 3; i < 2 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.9929289) < 1e-6);
+  }
+  for (size_t i = 2 * 3 * 3; i < 3 * 3 * 3; ++i) {
+    EXPECT_TRUE(std::fabs(var_[i] - 0.9910557) < 1e-6);
+  }
+}
+}  // namespace kernel
+}  // namespace mindspore