Add batch size and padding size.

PiperOrigin-RevId: 338356310
Change-Id: I33c6d9e10f8059deed85e5a763dd44ceff5683fe
This commit is contained in:
Tianrun Li 2020-10-21 15:48:07 -07:00 committed by TensorFlower Gardener
parent eb68be4a42
commit 0de96336c4
2 changed files with 17 additions and 1 deletions
tensorflow/core/profiler/utils

View File

@ -112,6 +112,10 @@ const HostEventTypeMap& GetHostEventTypeMap() {
// Batching related.
{"BatchingSessionRun", kBatchingSessionRun},
{"ProcessBatch", kProcessBatch},
{"ConcatInputTensors", kConcatInputTensors},
{"MergeInputTensors", kMergeInputTensors},
{"ScheduleWithoutSplit", kScheduleWithoutSplit},
{"ScheduleWithSplit", kScheduleWithSplit},
// JAX related.
{"LocalExecutable::ExecuteOnLocalDevices", kExecuteOnLocalDevices},
// GPU related.
@ -202,6 +206,10 @@ const StatTypeMap& GetStatTypeMap() {
{"memory_size", kDevCapMemorySize},
{"compute_cap_major", kDevCapComputeCapMajor},
{"compute_cap_minor", kDevCapComputeCapMinor},
// Batching related.
{"batch_size_after_padding", kBatchSizeAfterPadding},
{"padding_amount", kPaddingAmount},
{"batching_input_task_size", kBatchingInputTaskSize},
});
DCHECK_EQ(stat_type_map->size(), kNumStatTypes);
return *stat_type_map;

View File

@ -103,6 +103,10 @@ enum HostEventType {
// Batching related.
kBatchingSessionRun,
kProcessBatch,
kConcatInputTensors,
kMergeInputTensors,
kScheduleWithoutSplit,
kScheduleWithSplit,
// JAX related.
kExecuteOnLocalDevices,
// GPU related.
@ -191,7 +195,11 @@ enum StatType {
kDevCapMemorySize,
kDevCapComputeCapMajor,
kDevCapComputeCapMinor,
kLastStatType = kDevCapComputeCapMinor,
// Batching related.
kBatchSizeAfterPadding,
kPaddingAmount,
kBatchingInputTaskSize,
kLastStatType = kBatchingInputTaskSize,
};
inline std::string GpuPlaneName(int32 device_ordinal) {