[TF-XLA] Implement FtrlOptimizer

Change the TF documentation for the operation assigned to `linear` variable in ResourceApplyFtrl training_ops.

PiperOrigin-RevId: 158565492

tensorflow/compiler/tests/BUILD

@@ -183,6 +183,20 @@ tf_xla_py_test(
  ],
 )
 
+tf_xla_py_test(
+ name = "ftrl_test",
+ size = "small",
+ srcs = ["ftrl_test.py"],
+ deps = [
+ ":xla_test",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:framework_for_generated_wrappers",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:platform_test",
+ "//tensorflow/python:training",
+ ],
+)
+
 tf_xla_py_test(
  name = "function_test",
  size = "small",

tensorflow/compiler/tests/ftrl_test.py

@@ -0,0 +1,253 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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:https://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.
+# ==============================================================================
+"""Tests for Ftrl optimizer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.compiler.tests.xla_test import XLATestCase
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad
+from tensorflow.python.training import ftrl
+from tensorflow.python.training import gradient_descent
+
+
+class FtrlOptimizerTest(XLATestCase):
+
+ def initVariableAndGradient(self, dtype):
+ var0 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+ var1 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+ grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+ grads1 = constant_op.constant([0.02, 0.04], dtype=dtype)
+
+ return var0, var1, grads0, grads1
+
+ def equivAdagradTest_FtrlPart(self, steps, dtype):
+ var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype)
+ opt = ftrl.FtrlOptimizer(
+ 3.0,
+ learning_rate_power=-0.5, # using Adagrad learning rate
+ initial_accumulator_value=0.1,
+ l1_regularization_strength=0.0,
+ l2_regularization_strength=0.0)
+ ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+ variables.global_variables_initializer().run()
+ # Fetch params to validate initial values
+ self.assertAllClose([0.0, 0.0], var0.eval())
+ self.assertAllClose([0.0, 0.0], var1.eval())
+
+ # Run Ftrl for a few steps
+ for _ in range(steps):
+ ftrl_update.run()
+
+ return var0.eval(), var1.eval()
+
+ def equivAdagradTest_AdagradPart(self, steps, dtype):
+ var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype)
+ opt = adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1)
+ adagrad_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+ variables.global_variables_initializer().run()
+ # Fetch params to validate initial values
+ self.assertAllClose([0.0, 0.0], var0.eval())
+ self.assertAllClose([0.0, 0.0], var1.eval())
+
+ # Run Adagrad for a few steps
+ for _ in range(steps):
+ adagrad_update.run()
+
+ return var0.eval(), var1.eval()
+
+ def equivGradientDescentTest_FtrlPart(self, steps, dtype):
+ var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype)
+ opt = ftrl.FtrlOptimizer(
+ 3.0,
+ learning_rate_power=-0.0, # using Fixed learning rate
+ initial_accumulator_value=0.1,
+ l1_regularization_strength=0.0,
+ l2_regularization_strength=0.0)
+ ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+ variables.global_variables_initializer().run()
+ # Fetch params to validate initial values
+ self.assertAllClose([0.0, 0.0], var0.eval())
+ self.assertAllClose([0.0, 0.0], var1.eval())
+
+ # Run Ftrl for a few steps
+ for _ in range(steps):
+ ftrl_update.run()
+
+ return var0.eval(), var1.eval()
+
+ def equivGradientDescentTest_GradientDescentPart(self, steps, dtype):
+ var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype)
+ opt = gradient_descent.GradientDescentOptimizer(3.0, name="sgd")
+ sgd_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+ variables.global_variables_initializer().run()
+ # Fetch params to validate initial values
+ self.assertAllClose([0.0, 0.0], var0.eval())
+ self.assertAllClose([0.0, 0.0], var1.eval())
+
+ # Run GradientDescent for a few steps
+ for _ in range(steps):
+ sgd_update.run()
+
+ return var0.eval(), var1.eval()
+
+ def testFtrlwithoutRegularization(self):
+ for dtype in self.float_types:
+ with self.test_session(), self.test_scope():
+ var0 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+ var1 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+ grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+ grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+ opt = ftrl.FtrlOptimizer(
+ 3.0,
+ initial_accumulator_value=0.1,
+ l1_regularization_strength=0.0,
+ l2_regularization_strength=0.0)
+ ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+ variables.global_variables_initializer().run()
+ # Fetch params to validate initial values
+ self.assertAllClose([0.0, 0.0], var0.eval())
+ self.assertAllClose([0.0, 0.0], var1.eval())
+
+ # Run 3 steps FTRL
+ for _ in range(3):
+ ftrl_update.run()
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(
+ np.array([-2.60260963, -4.29698515]), var0.eval())
+ self.assertAllCloseAccordingToType(
+ np.array([-0.28432083, -0.56694895]), var1.eval())
+
+ def testFtrlwithoutRegularization2(self):
+ for dtype in self.float_types:
+ with self.test_session(), self.test_scope():
+ var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+ var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+ grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+ grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+ opt = ftrl.FtrlOptimizer(
+ 3.0,
+ initial_accumulator_value=0.1,
+ l1_regularization_strength=0.0,
+ l2_regularization_strength=0.0)
+ ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+ variables.global_variables_initializer().run()
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], var0.eval())
+ self.assertAllClose([4.0, 3.0], var1.eval())
+
+ # Run 3 steps FTRL
+ for _ in range(3):
+ ftrl_update.run()
+
+ # Validate updated params
+ self.assertAllClose(
+ np.array([-2.55607247, -3.98729396]), var0.eval(), 1e-5, 1e-5)
+ self.assertAllClose(
+ np.array([-0.28232238, -0.56096673]), var1.eval(), 1e-5, 1e-5)
+
+ def testFtrlWithL1(self):
+ for dtype in self.float_types:
+ with self.test_session(), self.test_scope():
+ var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+ var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+ grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+ grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+ opt = ftrl.FtrlOptimizer(
+ 3.0,
+ initial_accumulator_value=0.1,
+ l1_regularization_strength=0.001,
+ l2_regularization_strength=0.0)
+ ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+ variables.global_variables_initializer().run()
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], var0.eval())
+ self.assertAllClose([4.0, 3.0], var1.eval())
+
+ # Run 10 steps FTRL
+ for _ in range(10):
+ ftrl_update.run()
+
+ # Validate updated params
+ self.assertAllClose(np.array([-7.66718769, -10.91273689]), var0.eval())
+ self.assertAllClose(np.array([-0.93460727, -1.86147261]), var1.eval())
+
+ def testFtrlWithL1_L2(self):
+ for dtype in self.float_types:
+ with self.test_session(), self.test_scope():
+ var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+ var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+ grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+ grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+ opt = ftrl.FtrlOptimizer(
+ 3.0,
+ initial_accumulator_value=0.1,
+ l1_regularization_strength=0.001,
+ l2_regularization_strength=2.0)
+ ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+ variables.global_variables_initializer().run()
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], var0.eval())
+ self.assertAllClose([4.0, 3.0], var1.eval())
+
+ # Run 10 steps FTRL
+ for _ in range(10):
+ ftrl_update.run()
+
+ # Validate updated params
+ self.assertAllClose(np.array([-0.24059935, -0.46829352]), var0.eval())
+ self.assertAllClose(np.array([-0.02406147, -0.04830509]), var1.eval())
+
+ # When variables are intialized with Zero, FTRL-Proximal has two properties:
+ # 1. Without L1&L2 but with fixed learning rate, FTRL-Proximal is identical
+ # with GradientDescent.
+ # 2. Without L1&L2 but with adaptive learning rate, FTRL-Proximal is idential
+ # with Adagrad.
+ # So, basing on these two properties, we test if our implementation of
+ # FTRL-Proximal performs same updates as Adagrad or GradientDescent.
+ def testEquivAdagradwithoutRegularization(self):
+ steps = 5
+ for dtype in self.float_types:
+ with self.test_session(), self.test_scope():
+ val0, val1 = self.equivAdagradTest_FtrlPart(steps, dtype)
+ with self.test_session(), self.test_scope():
+ val2, val3 = self.equivAdagradTest_AdagradPart(steps, dtype)
+
+ self.assertAllClose(val0, val2)
+ self.assertAllClose(val1, val3)
+
+ def testEquivGradientDescentwithoutRegularization(self):
+ steps = 5
+ for dtype in self.float_types:
+ with self.test_session(), self.test_scope():
+ val0, val1 = self.equivGradientDescentTest_FtrlPart(steps, dtype)
+ with self.test_session(), self.test_scope():
+ val2, val3 = self.equivGradientDescentTest_GradientDescentPart(
+ steps, dtype)
+
+ self.assertAllClose(val0, val2)
+ self.assertAllClose(val1, val3)
+
+
+if __name__ == "__main__":
+ test.main()

tensorflow/compiler/tf2xla/kernels/training_ops.cc

@@ -364,5 +364,112 @@ class ResourceApplyRMSProp : public XlaOpKernel {
 };
 REGISTER_XLA_OP(Name("ResourceApplyRMSProp"), ResourceApplyRMSProp);
 
+class ResourceApplyFtrl : public XlaOpKernel {
+ public:
+ explicit ResourceApplyFtrl(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+ OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_));
+ }
+
+ void Compile(XlaOpKernelContext* ctx) override {
+ xla::ComputationBuilder* b = ctx->builder();
+
+ DataType var_type, accum_type, linear_type;
+ TensorShape var_shape, accum_shape, linear_shape;
+ OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &var_type, &var_shape));
+ OP_REQUIRES_OK(ctx,
+ ctx->GetVariableTypeAndShape(1, &accum_type, &accum_shape));
+ OP_REQUIRES_OK(
+ ctx, ctx->GetVariableTypeAndShape(2, &linear_type, &linear_shape));
+
+ OP_REQUIRES(
+ ctx,
+ dtype_ == var_type && dtype_ == accum_type && dtype_ == linear_type,
+ errors::InvalidArgument(
+ "Types of variable arguments to ResourceApplyFtrl must match: ",
+ DataTypeString(dtype_), " vs. ", DataTypeString(var_type), " and ",
+ DataTypeString(accum_type), " and ", DataTypeString(linear_type)));
+
+ OP_REQUIRES(ctx, var_shape.IsSameSize(accum_shape),
+ errors::InvalidArgument(
+ "var and accum do not have the same shape",
+ var_shape.DebugString(), " ", accum_shape.DebugString()));
+
+ OP_REQUIRES(ctx, var_shape.IsSameSize(linear_shape),
+ errors::InvalidArgument(
+ "var and linear do not have the same shape",
+ var_shape.DebugString(), " ", linear_shape.DebugString()));
+
+ TensorShape grad_shape = ctx->InputShape(3);
+ TensorShape lr_shape = ctx->InputShape(4);
+ TensorShape l1_shape = ctx->InputShape(5);
+ TensorShape l2_shape = ctx->InputShape(6);
+ TensorShape lr_power_shape = ctx->InputShape(7);
+
+ OP_REQUIRES(ctx, var_shape.IsSameSize(grad_shape),
+ errors::InvalidArgument(
+ "var and grad do not have the same shape",
+ var_shape.DebugString(), " ", grad_shape.DebugString()));
+
+ OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(lr_shape),
+ errors::InvalidArgument("lr is not a scalar: ",
+ lr_shape.DebugString()));
+
+ OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(l1_shape),
+ errors::InvalidArgument("l1 is not a scalar: ",
+ l1_shape.DebugString()));
+
+ OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(l2_shape),
+ errors::InvalidArgument("l2 is not a scalar: ",
+ l2_shape.DebugString()));
+
+ OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(lr_power_shape),
+ errors::InvalidArgument("lr_power is not a scalar: ",
+ lr_power_shape.DebugString()));
+
+ xla::ComputationDataHandle var, accum, linear;
+ OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &var));
+ OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, &accum));
+ OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, &linear));
+ xla::ComputationDataHandle grad = ctx->Input(3);
+ xla::ComputationDataHandle lr = ctx->Input(4);
+ xla::ComputationDataHandle l1 = ctx->Input(5);
+ xla::ComputationDataHandle l2 = ctx->Input(6);
+ xla::ComputationDataHandle lr_power = ctx->Input(7);
+
+ // new_accum = accum + grad * grad
+ // linear += grad - (new_accum^(-lr_power) - accum^(-lr_power)) / lr * var
+ // quadratic = (new_accum^(-lr_power) / lr) + 2 * l2
+ // var = (sign(linear) * l1 - linear) / quadratic if |linear| > l1 else 0.0
+ // accum = new_accum
+
+ xla::ComputationDataHandle zero_broadcast = b->Broadcast(
+ XlaHelpers::FloatLiteral(b, dtype_, 0.0), var_shape.dim_sizes());
+ xla::ComputationDataHandle two = XlaHelpers::FloatLiteral(b, dtype_, 2.0);
+
+ xla::ComputationDataHandle new_accum = b->Add(accum, b->Pow(grad, two));
+ xla::ComputationDataHandle new_accum_lr_pow =
+ b->Pow(new_accum, b->Neg(lr_power));
+ xla::ComputationDataHandle accum_lr_pow = b->Pow(accum, b->Neg(lr_power));
+ linear = b->Add(
+ linear,
+ b->Sub(grad, b->Mul(b->Div(b->Sub(new_accum_lr_pow, accum_lr_pow), lr),
+ var)));
+ xla::ComputationDataHandle quadratic =
+ b->Add(b->Div(new_accum_lr_pow, lr), b->Mul(two, l2));
+ xla::ComputationDataHandle pre_shrink =
+ b->Div(b->Sub(b->Mul(l1, b->Sign(linear)), linear), quadratic);
+ var = b->Select(b->Gt(b->Abs(linear), l1), pre_shrink, zero_broadcast);
+ accum = new_accum;
+
+ OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, dtype_, var));
+ OP_REQUIRES_OK(ctx, ctx->AssignVariable(1, dtype_, accum));
+ OP_REQUIRES_OK(ctx, ctx->AssignVariable(2, dtype_, linear));
+ }
+
+ private:
+ DataType dtype_;
+};
+REGISTER_XLA_OP(Name("ResourceApplyFtrl"), ResourceApplyFtrl);
+
 } // namespace
 } // namespace tensorflow

tensorflow/core/ops/training_ops.cc

@@ -925,7 +925,7 @@ REGISTER_OP("ResourceApplyFtrl")
 Update '*var' according to the Ftrl-proximal scheme.
 
 accum_new = accum + grad * grad
-linear += grad + (accum_new^(-lr_power) - accum^(-lr_power)) / lr * var
+linear += grad - (accum_new^(-lr_power) - accum^(-lr_power)) / lr * var
 quadratic = 1.0 / (accum_new^(lr_power) * lr) + 2 * l2
 var = (sign(linear) * l1 - linear) / quadratic if |linear| > l1 else 0.0
 accum = accum_new