Function CreateGraphDef

tensorflow/core/distributed_runtime/rpcbench_test.cc:124–155

Make a program with specified number of stages and "width" ops per stage.

// Make a program with specified number of stages and "width" ops per stage.
GraphDef CreateGraphDef(int num_stages, int width, int tensor_size,
                        bool use_multiple_devices, const Cluster* cluster) {
  CHECK_GE(cluster->devices.size(), width);

  using namespace ::tensorflow::ops;  // NOLINT(build/namespaces)

  Scope s = Scope::NewRootScope();

  // x is from the feed.
  Output x = Const(s.WithOpName("x"), 0.0f, {tensor_size, 1});

  // Create stages.
  std::vector<Output> last_stage;
  last_stage.push_back(x);
  for (int i = 0; i < num_stages; i++) {
    std::vector<Output> this_stage;
    for (int j = 0; j < width; j++) {
      Output combine = AddN(
          s.WithDevice(cluster->devices[use_multiple_devices ? j : 0].name()),
          last_stage);
      this_stage.push_back(combine);
    }
    last_stage = this_stage;
  }

  // Create output.
  /* Output y =*/AddN(s.WithOpName("y"), last_stage);

  GraphDef def;
  TF_CHECK_OK(s.ToGraphDef(&def));
  return def;
}

string DebugString(const Tensor& x, const Tensor& y, int tensor_size) {
  CHECK_EQ(x.NumElements(), tensor_size);