[tf.data] Add vlog to record how much time each optimization step takes.

PiperOrigin-RevId: 345557199
Change-Id: I54f8f549b10e98ac605905853505e828e9502f45
This commit is contained in:
Jay Shi 2020-12-03 16:11:08 -08:00 committed by TensorFlower Gardener
parent 77f04991aa
commit a0a365e408

View File

@ -219,8 +219,13 @@ class ModelDatasetOp::Dataset : public DatasetBase {
tf_shared_lock l(mu_);
model_input_time = SelfInputTime();
}
int64 optimization_start_us = EnvTime::NowMicros();
model_->Optimize(dataset()->algorithm_, cpu_budget_, ram_budget_,
/*model_input_time=*/0);
VLOG(2) << "Optimized for "
<< (EnvTime::NowMicros() - optimization_start_us) << " us.";
// Exponentially increase the period of running the optimization
// until a threshold is reached.
if (optimization_period_ms != kOptimizationPeriodThresholdMs) {