Method Run

tensorflow/cc/training/prefetch_runner.cc:86–142

}

void PrefetchRunner::Run(size_t index) {
  Status status;

  // initialize tensor buffer.
  status = sess_->Run({}, {}, {params_.resume_op()}, nullptr);
  if (status != Status::OK()) {
    DealWithCancelledUnexpectedError(status);
    return;
  }

  // get value generator
  auto feed_in_tensors = params_.named_feed_input_tensors();
  std::vector<std::string> val_consume_ops;
  std::vector<std::string> val_gen_tensors;
  val_consume_ops.reserve(feed_in_tensors.size());
  val_gen_tensors.reserve(feed_in_tensors.size());
  for (auto iter = feed_in_tensors.begin(); iter != feed_in_tensors.end(); ++iter) {
    val_consume_ops.emplace_back(iter->first);
    val_gen_tensors.emplace_back(iter->second);
  }

  std::vector<Tensor> gen_val;
  std::vector<std::pair<std::string, Tensor>> gen_inputs;
  gen_val.reserve(feed_in_tensors.size());
  gen_inputs.reserve(feed_in_tensors.size());
  while (true) {
    if (TF_PREDICT_FALSE(coord_ && coord_->ShouldStop())) {
      mutex_lock l (mu_);
      thread_nums_--;
      return;
    }

    // generate feed tensor
    if (!val_gen_tensors.empty()) {
      gen_val.clear();
      gen_inputs.clear();
      status = sess_->Run(params_.run_options(), {}, val_gen_tensors, {},
                         &gen_val, nullptr);
      if (TF_PREDICT_FALSE(status != Status::OK())) {
        if (!CheckRunErrorStatus(status))
          return;
      }

      for (size_t i = 0; i < val_gen_tensors.size(); i++) {
        gen_inputs.emplace_back(val_consume_ops[i], gen_val[i]);
      }
    }

    // run prefetch subgraph.
    status = sess_->Run(params_.run_options(), gen_inputs, {},
                       {params_.fetch_ops(index)}, nullptr, nullptr);
    if (TF_PREDICT_FALSE(status != Status::OK())) {
      if (!CheckRunErrorStatus(status))
        return;
    }
  }
}