-
Notifications
You must be signed in to change notification settings - Fork 5.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Implementing the Adamax optimizer operator #4538
Changes from 12 commits
8ff3a45
be99868
be2d958
abd6181
0b98f7d
1a4dda0
2e39197
9a1f9aa
17e1ea5
42b2166
7c921a8
0ae77e8
cdc46c7
af36e75
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,139 @@ | ||
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve. | ||
|
||
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 "paddle/operators/adamax_op.h" | ||
|
||
namespace paddle { | ||
namespace operators { | ||
|
||
class AdamaxOp : public framework::OperatorWithKernel { | ||
public: | ||
using framework::OperatorWithKernel::OperatorWithKernel; | ||
|
||
protected: | ||
void InferShape(framework::InferShapeContextBase *ctx) const override { | ||
PADDLE_ENFORCE(ctx->HasInput("Param"), | ||
"Input(Param) of AdamaxOp should not be null."); | ||
PADDLE_ENFORCE(ctx->HasInput("Grad"), | ||
"Input(Grad) of AdamaxOp should not be null."); | ||
PADDLE_ENFORCE(ctx->HasInput("Moment"), | ||
"Input(Moment) of AdamaxOp should not be null."); | ||
PADDLE_ENFORCE(ctx->HasInput("InfNorm"), | ||
"Input(InfNorm) of AdamaxOp should not be null."); | ||
PADDLE_ENFORCE(ctx->HasInput("LearningRate"), | ||
"Input(LearningRate) of AdamaxOp should not be null."); | ||
PADDLE_ENFORCE(ctx->HasInput("Beta1Pow"), | ||
"Input(Beta1Pow) of AdamaxOp should not be null."); | ||
|
||
PADDLE_ENFORCE(ctx->HasOutput("ParamOut"), | ||
"Output(ParamOut) of AdamaxOp should not be null."); | ||
PADDLE_ENFORCE(ctx->HasOutput("MomentOut"), | ||
"Output(MomentOut) of AdamaxOp should not be null."); | ||
PADDLE_ENFORCE(ctx->HasOutput("InfNormOut"), | ||
"Output(InfNormOut) of AdamaxOp should not be null."); | ||
PADDLE_ENFORCE(ctx->HasOutput("Beta1PowOut"), | ||
"Output(Beta1PowOut) of AdamaxOp should not be null."); | ||
|
||
auto lr_dims = ctx->GetInputDim("LearningRate"); | ||
PADDLE_ENFORCE_EQ(framework::product(lr_dims), 1, | ||
"Learning rate should have 1 dimension"); | ||
auto beta1_pow_dims = ctx->GetInputDim("Beta1Pow"); | ||
PADDLE_ENFORCE_EQ(framework::product(beta1_pow_dims), 1, | ||
"Beta1 power accumulator should have 1 dimension"); | ||
auto param_dims = ctx->GetInputDim("Param"); | ||
PADDLE_ENFORCE_EQ( | ||
param_dims, ctx->GetInputDim("Grad"), | ||
"Param and Grad input of AdamaxOp should have same dimension"); | ||
PADDLE_ENFORCE_EQ( | ||
param_dims, ctx->GetInputDim("Moment"), | ||
"Param and Moment input of AdamaxOp should have same dimension"); | ||
PADDLE_ENFORCE_EQ( | ||
param_dims, ctx->GetInputDim("InfNorm"), | ||
"Param and InfNorm input of AdamaxOp should have same dimension"); | ||
|
||
ctx->SetOutputDim("ParamOut", param_dims); | ||
ctx->SetOutputDim("MomentOut", param_dims); | ||
ctx->SetOutputDim("InfNormOut", param_dims); | ||
ctx->SetOutputDim("Beta1PowOut", beta1_pow_dims); | ||
} | ||
}; | ||
|
||
class AdamaxOpMaker : public framework::OpProtoAndCheckerMaker { | ||
public: | ||
AdamaxOpMaker(framework::OpProto *proto, framework::OpAttrChecker *op_checker) | ||
: OpProtoAndCheckerMaker(proto, op_checker) { | ||
AddInput("Param", "(Tensor) Input parameter"); | ||
AddInput("Grad", "(Tensor) Input gradient"); | ||
AddInput("LearningRate", "(Tensor) Learning rate"); | ||
AddInput("Moment", "(Tensor) First moment"); | ||
AddInput("InfNorm", | ||
"(Tensor) " | ||
"Input exponentially weighted infinity norm"); | ||
AddInput("Beta1Pow", "(Tensor) Input beta1 power accumulator"); | ||
|
||
AddOutput("ParamOut", "(Tensor) Output parameter"); | ||
AddOutput("MomentOut", "(Tensor) Output first moment"); | ||
AddOutput("InfNormOut", | ||
"(Tensor) " | ||
"Output exponentially weighted infinity norm"); | ||
AddOutput("Beta1PowOut", "(Tensor) Output beta1 power accumulator"); | ||
|
||
AddAttr<float>("beta1", | ||
"(float, default 0.9) " | ||
"Exponential decay rate for the " | ||
"1st moment estimates.") | ||
.SetDefault(0.9f); | ||
AddAttr<float>("beta2", | ||
"(float, default 0.999) " | ||
"exponential decay rate for the weighted " | ||
"infinity norm estimates.") | ||
.SetDefault(0.999f); | ||
AddAttr<float>("epsilon", | ||
"(float, default 1.0e-8) " | ||
"Constant for numerical stability") | ||
.SetDefault(1.0e-8f); | ||
AddComment(R"DOC( | ||
Adamax Updates Operator. | ||
|
||
This implements the Adamax optimizer from Section 7 of the Adam | ||
paper[1]. Adamax is a variant of the | ||
Adam algorithm based on the infinity norm. | ||
|
||
Adamax updates: | ||
|
||
moment_out = beta1 * moment + (1 - beta1) * grad | ||
inf_norm_out = max(beta2 * inf_norm + epsilon, abs(grad)) | ||
beta1_pow_out = beta1_pow * beta1 | ||
learning_rate_t = learning_rate/(1 - beta1_pow_out) | ||
param_out = param - learning_rate_t * moment_out/inf_norm_out | ||
|
||
The original paper does not have an epsilon attribute. | ||
However, it is added here for numerical stability | ||
by preventing divide by 0. | ||
|
||
References: | ||
[1] Adam: A Method for Stochastic Optimization | ||
(https://arxiv.org/abs/1412.6980) | ||
|
||
)DOC"); | ||
} | ||
}; | ||
|
||
} // namespace operators | ||
} // namespace paddle | ||
|
||
namespace ops = paddle::operators; | ||
REGISTER_OP_WITHOUT_GRADIENT(adamax, ops::AdamaxOp, ops::AdamaxOpMaker); | ||
REGISTER_OP_CPU_KERNEL(adamax, | ||
ops::AdamaxOpKernel<paddle::platform::CPUPlace, float>); |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,20 @@ | ||
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve. | ||
|
||
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. */ | ||
|
||
#define EIGEN_USE_GPU | ||
#include "paddle/operators/adamax_op.h" | ||
|
||
namespace ops = paddle::operators; | ||
REGISTER_OP_GPU_KERNEL(adamax, | ||
ops::AdamaxOpKernel<paddle::platform::GPUPlace, float>); |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,72 @@ | ||
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve. | ||
|
||
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. */ | ||
|
||
#pragma once | ||
#include "paddle/framework/eigen.h" | ||
#include "paddle/framework/op_registry.h" | ||
|
||
namespace paddle { | ||
namespace operators { | ||
|
||
template <typename Place, typename T> | ||
class AdamaxOpKernel : public framework::OpKernel<T> { | ||
public: | ||
void Compute(const framework::ExecutionContext& ctx) const override { | ||
auto param_out_tensor = ctx.Output<framework::Tensor>("ParamOut"); | ||
auto moment_out_tensor = ctx.Output<framework::Tensor>("MomentOut"); | ||
auto inf_norm_out_tensor = ctx.Output<framework::Tensor>("InfNormOut"); | ||
auto beta1_pow_out_tensor = ctx.Output<framework::Tensor>("Beta1PowOut"); | ||
|
||
param_out_tensor->mutable_data<T>(ctx.GetPlace()); | ||
moment_out_tensor->mutable_data<T>(ctx.GetPlace()); | ||
inf_norm_out_tensor->mutable_data<T>(ctx.GetPlace()); | ||
beta1_pow_out_tensor->mutable_data<T>(ctx.GetPlace()); | ||
|
||
float beta1 = ctx.Attr<float>("beta1"); | ||
float beta2 = ctx.Attr<float>("beta2"); | ||
float epsilon = ctx.Attr<float>("epsilon"); | ||
|
||
auto param = framework::EigenVector<T>::Flatten( | ||
*ctx.Input<framework::Tensor>("Param")); | ||
auto grad = framework::EigenVector<T>::Flatten( | ||
*ctx.Input<framework::Tensor>("Grad")); | ||
auto moment = framework::EigenVector<T>::Flatten( | ||
*ctx.Input<framework::Tensor>("Moment")); | ||
auto inf_norm = framework::EigenVector<T>::Flatten( | ||
*ctx.Input<framework::Tensor>("InfNorm")); | ||
auto lr = framework::EigenVector<T>::Flatten( | ||
*ctx.Input<framework::Tensor>("LearningRate")); | ||
auto beta1_pow = framework::EigenVector<T>::Flatten( | ||
*ctx.Input<framework::Tensor>("Beta1Pow")); | ||
auto param_out = framework::EigenVector<T>::Flatten(*param_out_tensor); | ||
auto moment_out = framework::EigenVector<T>::Flatten(*moment_out_tensor); | ||
auto inf_norm_out = | ||
framework::EigenVector<T>::Flatten(*inf_norm_out_tensor); | ||
auto beta1_pow_out = | ||
framework::EigenVector<T>::Flatten(*beta1_pow_out_tensor); | ||
auto place = ctx.GetEigenDevice<Place>(); | ||
|
||
moment_out.device(place) = beta1 * moment + (1 - beta1) * grad; | ||
inf_norm_out.device(place) = | ||
grad.abs().cwiseMax((beta2 * inf_norm) + epsilon); | ||
beta1_pow_out.device(place) = beta1_pow * beta1; | ||
auto lr_t = lr / (1 - beta1_pow_out); | ||
Eigen::DSizes<int, 1> m_dsize(moment_out_tensor->numel()); | ||
param_out.device(place) = | ||
param - lr_t.broadcast(m_dsize) * (moment_out / inf_norm_out); | ||
} | ||
}; | ||
|
||
} // namespace operators | ||
} // namespace paddle |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,96 @@ | ||
import unittest | ||
import numpy as np | ||
from op_test import OpTest | ||
|
||
|
||
class TestAdamaxOp1(OpTest): | ||
def setUp(self): | ||
self.op_type = "adamax" | ||
param = np.random.uniform(-1, 1, (102, 105)).astype("float32") | ||
grad = np.random.uniform(-1, 1, (102, 105)).astype("float32") | ||
moment = np.random.uniform(-1, 1, (102, 105)).astype("float32") | ||
# The infinity norm is positive | ||
inf_norm = np.random.random((102, 105)).astype("float32") | ||
|
||
learning_rate = 0.002 | ||
beta_1 = 0.9 | ||
beta_2 = 0.999 | ||
epsilon = 1e-8 | ||
beta_1_pow = beta_1**8 | ||
|
||
self.inputs = { | ||
'Param': param, | ||
'Grad': grad, | ||
'Moment': moment, | ||
'InfNorm': inf_norm, | ||
'LearningRate': np.array([learning_rate]).astype("float32"), | ||
'Beta1Pow': np.array([beta_1_pow]).astype("float32") | ||
} | ||
|
||
self.attrs = {'beta1': beta_1, 'beta2': beta_2, 'epsilon': epsilon} | ||
|
||
moment_out = beta_1 * moment + (1 - beta_1) * grad | ||
inf_norm_out = np.maximum(beta_2 * inf_norm + epsilon, np.abs(grad)) | ||
beta_1_pow_out = beta_1_pow * beta_1 | ||
lr_t = (learning_rate / (1 - beta_1_pow_out)) | ||
param_out = param - lr_t * np.divide(moment_out, inf_norm_out) | ||
|
||
self.outputs = { | ||
'ParamOut': param_out, | ||
'MomentOut': moment_out, | ||
'InfNormOut': inf_norm_out, | ||
'Beta1PowOut': beta_1_pow_out | ||
} | ||
|
||
def test_check_output(self): | ||
self.check_output() | ||
|
||
|
||
class TestAdamaxOp2(OpTest): | ||
'''Test Adamax Operator with default attributes | ||
''' | ||
|
||
def setUp(self): | ||
self.op_type = "adamax" | ||
param = np.random.uniform(-1, 1, (102, 105)).astype("float32") | ||
grad = np.random.uniform(-1, 1, (102, 105)).astype("float32") | ||
moment = np.random.uniform(-1, 1, (102, 105)).astype("float32") | ||
# The infinity norm is positive | ||
inf_norm = np.random.random((102, 105)).astype("float32") | ||
|
||
learning_rate = 0.002 | ||
beta_1 = 0.9 | ||
beta_2 = 0.999 | ||
epsilon = 1e-8 | ||
beta_1_pow = beta_1**8 | ||
|
||
self.inputs = { | ||
'Param': param, | ||
'Grad': grad, | ||
'Moment': moment, | ||
'InfNorm': inf_norm, | ||
'LearningRate': np.array([learning_rate]).astype("float32"), | ||
'Beta1Pow': np.array([beta_1_pow]).astype("float32") | ||
} | ||
|
||
self.attrs = {'beta1': beta_1, 'beta2': beta_2, 'epsilon': epsilon} | ||
|
||
moment_out = beta_1 * moment + (1 - beta_1) * grad | ||
inf_norm_out = np.maximum(beta_2 * inf_norm + epsilon, np.abs(grad)) | ||
beta_1_pow_out = beta_1_pow * beta_1 | ||
lr_t = (learning_rate / (1 - beta_1_pow_out)) | ||
param_out = param - lr_t * np.divide(moment_out, inf_norm_out) | ||
|
||
self.outputs = { | ||
'ParamOut': param_out, | ||
'MomentOut': moment_out, | ||
'InfNormOut': inf_norm_out, | ||
'Beta1PowOut': beta_1_pow_out | ||
} | ||
|
||
def test_check_output(self): | ||
self.check_output() | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The test of this kind of operator(Optimizer with state) should be more complex because we have accumulated state. The state will change when running, so the test code should run multiple times to check if the state is right. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fixed in af36e75 |
||
|
||
|
||
if __name__ == "__main__": | ||
unittest.main() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You provided two TestAdamaxOp functions and commented that the second one is for testing default attributes. I think it would be helpful to also add a comment for TestAdamaxOp1 explaining its purpose. Also, I didn't find any differences between the two test functions. If the first function is to test explicit attributes, you should change the attribute values.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you. I forgot to remove the attributes from the second one . Will change this.