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

[tf.data] Include information about element shapes and types in TraceMe metadata.

Browse Source 
PiperOrigin-RevId: 358258835
Change-Id: I2652e2fe0d10f1462ac90d27f7d09e071a2f31ed
This commit is contained in:
Jiri Simsa 2021-02-18 13:58:28 -08:00 committed by TensorFlower Gardener
parent f8a5d4e6cb
commit c9240bd8ed
2 changed files with 23 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
tensorflow/core/framework/dataset.cc
View File

@ -26,6 +26,7 @@ limitations under the License.
#include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/resource.h" #include "tensorflow/core/platform/resource.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/profiler/lib/traceme.h" #include "tensorflow/core/profiler/lib/traceme.h"
namespace tensorflow { namespace tensorflow {
@ -527,6 +528,26 @@ DatasetBaseIterator::DatasetBaseIterator(const BaseParams& params)
: params_(params) { : params_(params) {
params_.dataset->Ref(); params_.dataset->Ref();
VLOG(2) << prefix() << " constructor"; VLOG(2) << prefix() << " constructor";
traceme_metadata_ = strings::StrCat("id=", id_);
if (parent_) {
strings::StrAppend(&traceme_metadata_, ",parent_id=", parent_id_);
}
strings::StrAppend(&traceme_metadata_, ",shapes=");
auto& shapes = output_shapes();
for (int i = 0; i < shapes.size(); ++i) {
if (i > 0) {
strings::StrAppend(&traceme_metadata_, " ");
}
strings::StrAppend(&traceme_metadata_, shapes.at(i).DebugString());
}
strings::StrAppend(&traceme_metadata_, ",types=");
auto& types = output_dtypes();
for (int i = 0; i < types.size(); ++i) {
if (i > 0) {
strings::StrAppend(&traceme_metadata_, " ");
}
strings::StrAppend(&traceme_metadata_, DataTypeString(types.at(i)));
}
} }
DatasetBaseIterator::~DatasetBaseIterator() { DatasetBaseIterator::~DatasetBaseIterator() {
@ -535,11 +556,7 @@ DatasetBaseIterator::~DatasetBaseIterator() {
} }
string DatasetBaseIterator::BuildTraceMeName() { string DatasetBaseIterator::BuildTraceMeName() {
string result = strings::StrCat(params_.prefix, "#id=", id_); string result = strings::StrCat(params_.prefix, "#", traceme_metadata_);
if (parent_) {
strings::StrAppend(&result, ",parent_id=", parent_id_);
}
TraceMeMetadata metadata = GetTraceMeMetadata(); TraceMeMetadata metadata = GetTraceMeMetadata();
for (const auto& pair : metadata) { for (const auto& pair : metadata) {
strings::StrAppend(&result, ",", pair.first, "=", pair.second); strings::StrAppend(&result, ",", pair.first, "=", pair.second);

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

@ -1094,6 +1094,7 @@ class DatasetBaseIterator : public IteratorBase {
return model && model->collect_resource_usage() && node_; return model && model->collect_resource_usage() && node_;
} }
string traceme_metadata_;
BaseParams params_; BaseParams params_;
}; };