Caffe2 - C++ API
A deep learning, cross platform ML framework
adam_op.cc
1 
17 #include "adam_op.h"
18 
19 namespace caffe2 {
20 
21 REGISTER_CPU_OPERATOR(Adam, AdamOp<float, CPUContext>);
22 OPERATOR_SCHEMA(Adam)
23  .NumInputs(6)
24  .NumOutputs(3)
25  .AllowInplace({{0, 0}, {1, 1}, {2, 2}})
26  .SetDoc(R"DOC(
27 
28 Computes the Adam update (https://arxiv.org/abs/1412.6980) for an
29 input gradient and momentum parameters. Concretely, given inputs
30 (param, m1, m2, grad, lr, iters),
31 
32  t = iters + 1
33  corrected_local_rate = lr * sqrt(1 - power(beta2, t)) /
34  (1 - power(beta1, t))
35  m1_o = (beta1 * m1) + (1 - beta1) * grad
36  m2_o = (beta2 * m2) + (1 - beta2) * np.square(grad)
37  grad_o = corrected_local_rate * m1_o / \
38  (sqrt(m2_o) + epsilon)
39  param_o = param + grad_o
40 
41 and returns (param_o, m1_o, m2_o)
42 
43 )DOC")
44  .Input(0, "param", "Parameters to be updated")
45  .Input(1, "moment_1", "First moment history")
46  .Input(2, "moment_2", "Second moment history")
47  .Input(3, "grad", "Gradient computed")
48  .Input(4, "lr", "learning rate")
49  .Input(5, "iter", "iteration number")
50  .Output(0, "output_param", "Updated parameters")
51  .Output(1, "output_moment_1", "Updated first moment")
52  .Output(2, "output_moment_2", "Updated second moment")
53  .Arg("beta1", "Default 0.9")
54  .Arg("beta2", "Default 0.999")
55  .Arg("epsilon", "Default 1e-5");
56 
57 REGISTER_CPU_OPERATOR(SparseAdam, SparseAdamOp<float, CPUContext>);
58 OPERATOR_SCHEMA(SparseAdam)
59  .NumInputs(7)
60  .NumOutputs(3)
61  .EnforceInplace({{0, 0}, {1, 1}, {2, 2}})
62  .SetDoc(R"DOC(
63 
64 Computes the Adam Update for the sparse case.
65 Given inputs (param, moment1, moment2, indices, grad, lr, iter), runs the dense
66 Adam on (param, moment1[indices], momemnt2[indices], lr, iter) and returns
67 (new_param, new_moment1, new_moment2) as in dense case
68 
69 )DOC")
70  .Input(0, "param", "Parameters to be updated")
71  .Input(1, "moment_1", "First moment history")
72  .Input(2, "moment_2", "Second moment history")
73  .Input(3, "indices", "Sparse indices")
74  .Input(4, "grad", "Gradient computed")
75  .Input(5, "lr", "learning rate")
76  .Input(6, "iter", "iteration number")
77  .Output(0, "output_param", "Updated parameters")
78  .Output(1, "output_moment_1", "Updated first moment")
79  .Output(2, "output_moment_2", "Updated second moment")
80  .Arg("beta1", "Default 0.9")
81  .Arg("beta2", "Default 0.999")
82  .Arg("epsilon", "Default 1e-5");
83 
84 REGISTER_CPU_OPERATOR(
85  RowWiseSparseAdam,
86  RowWiseSparseAdamOp<float, CPUContext>);
87 OPERATOR_SCHEMA(RowWiseSparseAdam)
88  .NumInputs(7)
89  .NumOutputs(3)
90  .EnforceInplace({{0, 0}, {1, 1}, {2, 2}})
91  .SetDoc(R"DOC(
92 
93 Computes a modified Adam Update for the sparse case.
94 Given inputs (param, moment1, moment2, indices, grad, lr, iter), runs the
95 Adam update on (param, moment1[indices], moment2[indices], lr, iter) and returns
96 (new_param, new_moment1, new_moment2), where moment2 is a 1D tensor
97 with length equal to the number of rows in param:
98 shape(moment2) == shape(param)[0]. Each element of moment2 is
99 applied to an entire row of param, and the new moment2 values are
100 calculated by averaging across the row.
101 
102 )DOC")
103  .Input(0, "param", "Parameters to be updated")
104  .Input(1, "moment_1", "First moment history")
105  .Input(2, "moment_2", "Second moment history")
106  .Input(3, "indices", "Sparse indices")
107  .Input(4, "grad", "Gradient computed")
108  .Input(5, "lr", "learning rate")
109  .Input(6, "iter", "iteration number")
110  .Output(0, "output_param", "Updated parameters")
111  .Output(1, "output_moment_1", "Updated first moment")
112  .Output(2, "output_moment_2", "Updated second moment")
113  .Arg("beta1", "Default 0.9")
114  .Arg("beta2", "Default 0.999")
115  .Arg("epsilon", "Default 1e-5");
116 
117 SHOULD_NOT_DO_GRADIENT(Adam);
118 SHOULD_NOT_DO_GRADIENT(SparseAdam);
119 SHOULD_NOT_DO_GRADIENT(RowWiseSparseAdam);
120 }
Copyright (c) 2016-present, Facebook, Inc.