diff options
Diffstat (limited to 'tensorflow/core/util/stats_calculator.cc')
| -rw-r--r-- | tensorflow/core/util/stats_calculator.cc | 21 |
1 files changed, 18 insertions, 3 deletions
diff --git a/tensorflow/core/util/stats_calculator.cc b/tensorflow/core/util/stats_calculator.cc index c4befbdb84..eb07754650 100644 --- a/tensorflow/core/util/stats_calculator.cc +++ b/tensorflow/core/util/stats_calculator.cc @@ -272,9 +272,24 @@ std::string StatsCalculator::GetOutputString() const { return stream.str(); } -void StatsCalculator::UpdateDetails( - const std::map<std::string, Detail>& details) { - details_.insert(details.begin(), details.end()); +void StatsCalculator::AddNodeStats(const std::string& name, + const std::string& type, int64_t run_order, + int64_t start_us, int64_t rel_end_us, + int64_t mem_used) { + Detail* detail = nullptr; + if (details_.find(name) == details_.end()) { + details_.insert({name, {}}); + detail = &details_.at(name); + detail->type = type; + detail->name = name; + detail->run_order = run_order; + } else { + detail = &details_.at(name); + } + detail->start_us.UpdateStat(start_us); + detail->rel_end_us.UpdateStat(rel_end_us); + detail->mem_used.UpdateStat(mem_used); + detail->times_called++; } } // namespace tensorflow |