Move SDCA optimizer kernels to core.

See: tensorflow/contrib/linear_optimizer/kernels/g3doc/readme.md
Change: 134452636

1 files changed, 241 insertions, 0 deletions

diff --git a/tensorflow/core/kernels/sdca_ops_test.cc b/tensorflow/core/kernels/sdca_ops_test.cc
new file mode 100644
index 0000000000..dcdc9ad1cf
--- /dev/null
+++ b/tensorflow/core/kernels/sdca_ops_test.cc
@@ -0,0 +1,241 @@
+/* Copyright 2016 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://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 "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/lib/random/random.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+
+namespace {
+
+const SessionOptions* GetSingleThreadedOptions() {
+  static const SessionOptions* const kSessionOptions = []() {
+    SessionOptions* const result = new SessionOptions();
+    result->config.set_intra_op_parallelism_threads(1);
+    result->config.set_inter_op_parallelism_threads(1);
+    result->config.add_session_inter_op_thread_pool()->set_num_threads(1);
+    return result;
+  }();
+  return kSessionOptions;
+}
+
+const SessionOptions* GetMultiThreadedOptions() {
+  static const SessionOptions* const kSessionOptions = []() {
+    SessionOptions* const result = new SessionOptions();
+    result->config.set_intra_op_parallelism_threads(0);  // Auto-configured.
+    result->config.set_inter_op_parallelism_threads(0);  // Auto-configured.
+    result->config.add_session_inter_op_thread_pool()->set_num_threads(
+        0);  // Auto-configured.
+    return result;
+  }();
+  return kSessionOptions;
+}
+
+Node* Var(Graph* const g, const int n) {
+  return test::graph::Var(g, DT_FLOAT, TensorShape({n}));
+}
+
+// Returns a vector of size 'nodes' with each node being of size 'node_size'.
+std::vector<Node*> VarVector(Graph* const g, const int nodes,
+                             const int node_size) {
+  std::vector<Node*> result;
+  for (int i = 0; i < nodes; ++i) {
+    result.push_back(Var(g, node_size));
+  }
+  return result;
+}
+
+Node* Zeros(Graph* const g, const TensorShape& shape) {
+  Tensor data(DT_FLOAT, shape);
+  data.flat<float>().setZero();
+  return test::graph::Constant(g, data);
+}
+
+Node* Zeros(Graph* const g, const int n) { return Zeros(g, TensorShape({n})); }
+
+Node* Ones(Graph* const g, const int n) {
+  Tensor data(DT_FLOAT, TensorShape({n}));
+  test::FillFn<float>(&data, [](const int i) { return 1.0f; });
+  return test::graph::Constant(g, data);
+}
+
+Node* SparseExampleIndices(Graph* const g, const int sparse_features_per_group,
+                           const int num_examples) {
+  const int x_size = num_examples * 4;
+  Tensor data(DT_INT64, TensorShape({x_size}));
+  test::FillFn<int64>(&data, [&](const int i) { return i / 4; });
+  return test::graph::Constant(g, data);
+}
+
+Node* SparseFeatureIndices(Graph* const g, const int sparse_features_per_group,
+                           const int num_examples) {
+  const int x_size = num_examples * 4;
+  Tensor data(DT_INT64, TensorShape({x_size}));
+  test::FillFn<int64>(
+      &data, [&](const int i) { return i % sparse_features_per_group; });
+  return test::graph::Constant(g, data);
+}
+
+Node* RandomZeroOrOne(Graph* const g, const int n) {
+  Tensor data(DT_FLOAT, TensorShape({n}));
+  test::FillFn<float>(&data, [](const int i) {
+    // Fill with 0.0 or 1.0 at random.
+    return (random::New64() % 2) == 0 ? 0.0f : 1.0f;
+  });
+  return test::graph::Constant(g, data);
+}
+
+Node* RandomZeroOrOneMatrix(Graph* const g, const int n, int d) {
+  Tensor data(DT_FLOAT, TensorShape({n, d}));
+  test::FillFn<float>(&data, [](const int i) {
+    // Fill with 0.0 or 1.0 at random.
+    return (random::New64() % 2) == 0 ? 0.0f : 1.0f;
+  });
+  return test::graph::Constant(g, data);
+}
+
+void GetGraphs(const int32 num_examples, const int32 num_sparse_feature_groups,
+               const int32 sparse_features_per_group,
+               const int32 num_dense_feature_groups,
+               const int32 dense_features_per_group, Graph** const init_g,
+               Graph** train_g) {
+  {
+    // Build initialization graph
+    Graph* g = new Graph(OpRegistry::Global());
+
+    // These nodes have to be created first, and in the same way as the
+    // nodes in the graph below.
+    std::vector<Node*> sparse_weight_nodes =
+        VarVector(g, num_sparse_feature_groups, sparse_features_per_group);
+    std::vector<Node*> dense_weight_nodes =
+        VarVector(g, num_dense_feature_groups, dense_features_per_group);
+    Node* const multi_zero = Zeros(g, sparse_features_per_group);
+    for (Node* n : sparse_weight_nodes) {
+      test::graph::Assign(g, n, multi_zero);
+    }
+    Node* const zero = Zeros(g, dense_features_per_group);
+    for (Node* n : dense_weight_nodes) {
+      test::graph::Assign(g, n, zero);
+    }
+
+    *init_g = g;
+  }
+
+  {
+    // Build execution graph
+    Graph* g = new Graph(OpRegistry::Global());
+
+    // These nodes have to be created first, and in the same way as the
+    // nodes in the graph above.
+    std::vector<Node*> sparse_weight_nodes =
+        VarVector(g, num_sparse_feature_groups, sparse_features_per_group);
+    std::vector<Node*> dense_weight_nodes =
+        VarVector(g, num_dense_feature_groups, dense_features_per_group);
+
+    std::vector<NodeBuilder::NodeOut> sparse_weights;
+    for (Node* n : sparse_weight_nodes) {
+      sparse_weights.push_back(NodeBuilder::NodeOut(n));
+    }
+    std::vector<NodeBuilder::NodeOut> dense_weights;
+    for (Node* n : dense_weight_nodes) {
+      dense_weights.push_back(NodeBuilder::NodeOut(n));
+    }
+
+    std::vector<NodeBuilder::NodeOut> sparse_example_indices;
+    std::vector<NodeBuilder::NodeOut> sparse_feature_indices;
+    std::vector<NodeBuilder::NodeOut> sparse_values;
+    for (int i = 0; i < num_sparse_feature_groups; ++i) {
+      sparse_example_indices.push_back(NodeBuilder::NodeOut(
+          SparseExampleIndices(g, sparse_features_per_group, num_examples)));
+    }
+    for (int i = 0; i < num_sparse_feature_groups; ++i) {
+      sparse_feature_indices.push_back(NodeBuilder::NodeOut(
+          SparseFeatureIndices(g, sparse_features_per_group, num_examples)));
+    }
+    for (int i = 0; i < num_sparse_feature_groups; ++i) {
+      sparse_values.push_back(
+          NodeBuilder::NodeOut(RandomZeroOrOne(g, num_examples * 4)));
+    }
+
+    std::vector<NodeBuilder::NodeOut> dense_features;
+    for (int i = 0; i < num_dense_feature_groups; ++i) {
+      dense_features.push_back(NodeBuilder::NodeOut(
+          RandomZeroOrOneMatrix(g, num_examples, dense_features_per_group)));
+    }
+
+    Node* const weights = Ones(g, num_examples);
+    Node* const labels = RandomZeroOrOne(g, num_examples);
+    Node* const example_state_data = Zeros(g, TensorShape({num_examples, 4}));
+
+    Node* sdca = nullptr;
+    TF_CHECK_OK(
+        NodeBuilder(g->NewName("sdca"), "DistributedSdcaLargeBatchSolver")
+            .Attr("loss_type", "logistic_loss")
+            .Attr("num_sparse_features", num_sparse_feature_groups)
+            .Attr("num_sparse_features_with_values", num_sparse_feature_groups)
+            .Attr("num_dense_features", num_dense_feature_groups)
+            .Attr("l1", 0.0)
+            .Attr("l2", 1.0)
+            .Attr("num_loss_partitions", 1)
+            .Attr("num_inner_iterations", 2)
+            .Input(sparse_example_indices)
+            .Input(sparse_feature_indices)
+            .Input(sparse_values)
+            .Input(dense_features)
+            .Input(weights)
+            .Input(labels)
+            .Input(sparse_weights)
+            .Input(dense_weights)
+            .Input(example_state_data)
+            .Finalize(g, &sdca));
+
+    *train_g = g;
+  }
+}
+
+void BM_SDCA(const int iters, const int num_examples) {
+  testing::StopTiming();
+  Graph* init = nullptr;
+  Graph* train = nullptr;
+  GetGraphs(num_examples, 20 /* sparse feature groups */,
+            5 /* sparse features per group */, 1 /* dense feature groups*/,
+            20 /* dense features per group */, &init, &train);
+  testing::StartTiming();
+  test::Benchmark("cpu", train, GetSingleThreadedOptions(), init).Run(iters);
+}
+
+void BM_SDCA_LARGE_SPARSE(const int iters, const int num_examples) {
+  testing::StopTiming();
+  Graph* init = nullptr;
+  Graph* train = nullptr;
+  GetGraphs(num_examples, 65 /* sparse feature groups */,
+            1e6 /* sparse features per group */, 0 /* dense feature groups*/,
+            0 /* dense features per group */, &init, &train);
+  testing::StartTiming();
+  test::Benchmark("cpu", train, GetMultiThreadedOptions(), init).Run(iters);
+}
+
+}  // namespace
+
+BENCHMARK(BM_SDCA)->Arg(128)->Arg(256)->Arg(512)->Arg(1024);
+BENCHMARK(BM_SDCA_LARGE_SPARSE)->Arg(128)->Arg(256)->Arg(512)->Arg(1024);
+
+}  // namespace tensorflow