experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[tf.data] Add vlogging for time spent restoring each stage of a dataset iterator.

Browse Source 
PiperOrigin-RevId: 351396158
Change-Id: I71ff55afa330ec01006cdb2d0b09de1dfba7454f
This commit is contained in:
Andrew Audibert 2021-01-12 10:16:18 -08:00 committed by TensorFlower Gardener
parent dce4796033
commit 62acd0bf19
1 changed files with 20 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
tensorflow/core/framework/dataset.h
View File

@ -665,7 +665,11 @@ class IteratorBase {
// Saves the state of this iterator. // Saves the state of this iterator.
virtual Status Save(SerializationContext* ctx, IteratorStateWriter* writer) { virtual Status Save(SerializationContext* ctx, IteratorStateWriter* writer) {
return SaveInternal(ctx, writer); int64 start_us = EnvTime::NowMicros();
TF_RETURN_IF_ERROR(SaveInternal(ctx, writer));
VLOG(1) << "Saved " << prefix() << " in "
<< (EnvTime::NowMicros() - start_us) << "us";
return Status::OK();
} }
protected: protected:
@ -675,21 +679,33 @@ class IteratorBase {
// Restores the state of this iterator. // Restores the state of this iterator.
Status Restore(IteratorContext* ctx, IteratorStateReader* reader) { Status Restore(IteratorContext* ctx, IteratorStateReader* reader) {
return RestoreInternal(ctx, reader); int64 start_us = EnvTime::NowMicros();
TF_RETURN_IF_ERROR(RestoreInternal(ctx, reader));
VLOG(1) << "Restored " << prefix() << " in "
<< (EnvTime::NowMicros() - start_us) << "us";
return Status::OK();
} }
// This is needed so that sub-classes of IteratorBase can call // This is needed so that sub-classes of IteratorBase can call
// `SaveInternal` on their input iterators. // `SaveInternal` on their input iterators.
Status SaveInput(SerializationContext* ctx, IteratorStateWriter* writer, Status SaveInput(SerializationContext* ctx, IteratorStateWriter* writer,
const std::unique_ptr<IteratorBase>& input) { const std::unique_ptr<IteratorBase>& input) {
return input->SaveInternal(ctx, writer); int64 start_us = EnvTime::NowMicros();
TF_RETURN_IF_ERROR(input->SaveInternal(ctx, writer));
VLOG(2) << "Saved " << input->prefix() << " in "
<< (EnvTime::NowMicros() - start_us) << "us";
return Status::OK();
} }
// This is needed so that sub-classes of IteratorBase can call // This is needed so that sub-classes of IteratorBase can call
// `RestoreInternal` on their input iterators. // `RestoreInternal` on their input iterators.
Status RestoreInput(IteratorContext* ctx, IteratorStateReader* reader, Status RestoreInput(IteratorContext* ctx, IteratorStateReader* reader,
const std::unique_ptr<IteratorBase>& input) { const std::unique_ptr<IteratorBase>& input) {
return input->RestoreInternal(ctx, reader); int64 start_us = EnvTime::NowMicros();
TF_RETURN_IF_ERROR(input->RestoreInternal(ctx, reader));
VLOG(2) << "Restored " << input->prefix() << " in "
<< (EnvTime::NowMicros() - start_us) << "us";
return Status::OK();
} }
// Saves the state of this iterator. // Saves the state of this iterator.