diff options
| author |  Yuefeng Zhou <yuefengz@google.com> | 2017-01-20 10:59:37 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-01-20 11:06:19 -0800 |
| commit | 448c0dc61c40a631bf5f1a9ec1327591a28e63e6 (patch) | |
| tree | c1aaa2cf06913f6e1c653898146b8ca555263203 /tensorflow | |
| parent | 8499dc1ee56b40a6e1f16ccbf3073a12523a0f58 (diff) | |
Fix the bug that when cpu_allocated is called before session initialization memory cannot be tracked.
Change: 145103741
Diffstat (limited to 'tensorflow')
| -rw-r--r-- | tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc | 10 | ||||
| -rw-r--r-- | tensorflow/core/framework/allocator.cc | 3 |
2 files changed, 5 insertions, 8 deletions
diff --git a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc index 3c0989c587..abc8970216 100644 --- a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc +++ b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc @@ -45,8 +45,6 @@ namespace tensorflow { namespace { TEST(DirectSessionWithTrackingAllocTest, CostModelTest) { - EnableCPUAllocatorFullStats(true); - Graph graph(OpRegistry::Global()); Tensor a_tensor(DT_FLOAT, TensorShape({2, 2})); @@ -100,9 +98,9 @@ TEST(DirectSessionWithTrackingAllocTest, CostModelTest) { EXPECT_EQ(2, shape.dim(0).size()); EXPECT_EQ(1, shape.dim(1).size()); if (node->name() == y->name()) { - EXPECT_EQ(5, cm->AllocationId(node, 0)); + EXPECT_EQ(3, cm->AllocationId(node, 0)); } else { - EXPECT_EQ(6, cm->AllocationId(node, 0)); + EXPECT_EQ(4, cm->AllocationId(node, 0)); } } EXPECT_LE(0, cm->MaxExecutionTime(node)); @@ -115,8 +113,6 @@ TEST(DirectSessionWithTrackingAllocTest, CostModelTest) { } static void TestHWAccelerator(bool enableHWTrace) { - EnableCPUAllocatorFullStats(true); - Graph graph(OpRegistry::Global()); Tensor a_tensor(DT_FLOAT, TensorShape({2, 2})); @@ -199,8 +195,6 @@ TEST(DirectSessionWithTrackingAllocTest, CostModelWithHardwareStats) { } TEST(DirectSessionWithTrackingAllocTest, CostGraph) { - EnableCPUAllocatorFullStats(true); - Graph graph(OpRegistry::Global()); Tensor a_tensor(DT_FLOAT, TensorShape({2, 2})); diff --git a/tensorflow/core/framework/allocator.cc b/tensorflow/core/framework/allocator.cc index 812ce4bfe7..ff31ad965b 100644 --- a/tensorflow/core/framework/allocator.cc +++ b/tensorflow/core/framework/allocator.cc @@ -120,6 +120,9 @@ Allocator* MakeCpuAllocator() { Allocator* cpu_allocator() { static Allocator* cpu_alloc = MakeCpuAllocator(); + if (cpu_allocator_collect_full_stats && !cpu_alloc->TracksAllocationSizes()) { + cpu_alloc = new TrackingAllocator(cpu_alloc, true); + } return cpu_alloc; } |