Function BM_LargeTensorWrite

tensorflow/core/kernels/save_op_test.cc:664–692

// Benchmark-related code below.

static void BM_LargeTensorWrite(int iters, int num_elements) {
  testing::StopTiming();

  // 4 * num_elements bytes total , since sizeof(float) == 4.
  Tensor tensor(DT_FLOAT, TensorShape({num_elements}));
  tensor.flat<float>().setZero();

  // Builds the graph.
  const tstring temp_filename =
      io::JoinPath(testing::TmpDir(), "benchmark_checkpoint");
  auto root = Scope::NewRootScope().ExitOnError();
  const tstring tensor_name = "my_tensor";
  ops::Save(root, temp_filename, {tensor_name}, {{tensor}});

  // Disables optimizations.
  SessionOptions session_options;
  session_options.config.mutable_graph_options()
      ->mutable_optimizer_options()
      ->set_opt_level(tensorflow::OptimizerOptions::L0);

  TF_CHECK_OK(root.status());
  Graph* g = new Graph(OpRegistry::Global());
  TF_CHECK_OK(root.ToGraph(g));
  VLOG(1) << "Save op's output path: " << temp_filename;
  VLOG(1) << "# nodes in Graph: " << g->num_nodes();

  testing::StartTiming();
  test::Benchmark("cpu", g, &session_options).Run(iters);
}
BENCHMARK(BM_LargeTensorWrite)->Arg((1 << 30) / 4 /* 1GB float tensor */);

}  // namespace