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/mutex.h"
#include "tensorflow/core/platform/resource.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/profiler/lib/traceme.h"
namespace tensorflow {
@ -527,6 +528,26 @@ DatasetBaseIterator::DatasetBaseIterator(const BaseParams& params)
: params_(params) {
params_.dataset->Ref();
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() {
@ -535,11 +556,7 @@ DatasetBaseIterator::~DatasetBaseIterator() {
}
string DatasetBaseIterator::BuildTraceMeName() {
string result = strings::StrCat(params_.prefix, "#id=", id_);
if (parent_) {
strings::StrAppend(&result, ",parent_id=", parent_id_);
}
string result = strings::StrCat(params_.prefix, "#", traceme_metadata_);
TraceMeMetadata metadata = GetTraceMeMetadata();
for (const auto& pair : metadata) {
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_;
}
string traceme_metadata_;
BaseParams params_;
};