Function DoCompute

tensorflow/core/kernels/sdca_ops.cc:130–231

TODO(shengx): The helper classes/methods are changed to support multiclass SDCA, which lead to changes within this function. Need to revisit the convergence once the multiclass SDCA is in.

// SDCA, which lead to changes within this function. Need to revisit the
// convergence once the multiclass SDCA is in.
void DoCompute(const ComputeOptions& options, OpKernelContext* const context) {
  ModelWeights model_weights;
  OP_REQUIRES_OK(context, model_weights.Initialize(context));

  Examples examples;
  OP_REQUIRES_OK(
      context,
      examples.Initialize(context, model_weights, options.num_sparse_features,
                          options.num_sparse_features_with_values,
                          options.num_dense_features));

  const Tensor* example_state_data_t;
  OP_REQUIRES_OK(context,
                 context->input("example_state_data", &example_state_data_t));
  TensorShape expected_example_state_shape({examples.num_examples(), 4});
  OP_REQUIRES(context,
              example_state_data_t->shape() == expected_example_state_shape,
              errors::InvalidArgument(
                  "Expected shape ", expected_example_state_shape.DebugString(),
                  " for example_state_data, got ",
                  example_state_data_t->shape().DebugString()));

  Tensor mutable_example_state_data_t(*example_state_data_t);
  auto example_state_data = mutable_example_state_data_t.matrix<float>();
  OP_REQUIRES_OK(context, context->set_output("out_example_state_data",
                                              mutable_example_state_data_t));

  if (options.adaptive) {
    OP_REQUIRES_OK(context,
                   examples.SampleAdaptiveProbabilities(
                       options.num_loss_partitions, options.regularizations,
                       model_weights, example_state_data, options.loss_updater,
                       /*num_weight_vectors =*/1));
  } else {
    examples.RandomShuffle();
  }
  struct {
    mutex mu;
    Status value GUARDED_BY(mu);
  } train_step_status;
  std::atomic<std::int64_t> atomic_index(-1);
  auto train_step = [&](const int64 begin, const int64 end) {
    // The static_cast here is safe since begin and end can be at most
    // num_examples which is an int.
    for (int id = static_cast<int>(begin); id < end; ++id) {
      const int64 example_index = examples.sampled_index(++atomic_index);
      const Example& example = examples.example(example_index);
      const float dual = example_state_data(example_index, 0);
      const float example_weight = example.example_weight();
      float example_label = example.example_label();
      const Status conversion_status =
          options.loss_updater->ConvertLabel(&example_label);
      if (!conversion_status.ok()) {
        mutex_lock l(train_step_status.mu);
        train_step_status.value = conversion_status;
        // Return from this worker thread - the calling thread is
        // responsible for checking context status and returning on error.
        return;