From 413543eb8f9ff6939eee457974034afcb3e08718 Mon Sep 17 00:00:00 2001
From: Wei Liu <liuwei12@baidu.com>
Date: Fri, 18 Jan 2019 09:52:36 +0800
Subject: [PATCH] print peak memory usage

---
 paddle/fluid/memory/detail/system_allocator.cc | 5 +++++
 paddle/fluid/memory/detail/system_allocator.h  | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/paddle/fluid/memory/detail/system_allocator.cc b/paddle/fluid/memory/detail/system_allocator.cc
index 3e8fb83e9d5..14dcaf756f7 100644
--- a/paddle/fluid/memory/detail/system_allocator.cc
+++ b/paddle/fluid/memory/detail/system_allocator.cc
@@ -117,6 +117,11 @@ void* GPUAllocator::Alloc(size_t* index, size_t size) {
   if (result == cudaSuccess) {
     *index = 0;
     gpu_alloc_size_ += size;
+    if (gpu_alloc_size_ > s_memoryMap[gpu_id_]) {
+        s_memoryMap[gpu_id_] = gpu_alloc_size_;
+        VLOG(3) << "device: " << gpu_id_
+          << " maximum memory size : " <<(gpu_alloc_size_ >> 20) << " MiB";
+    }
     return p;
   } else {
     LOG(WARNING)
diff --git a/paddle/fluid/memory/detail/system_allocator.h b/paddle/fluid/memory/detail/system_allocator.h
index a0386a2dad1..1ac1df6de76 100644
--- a/paddle/fluid/memory/detail/system_allocator.h
+++ b/paddle/fluid/memory/detail/system_allocator.h
@@ -15,6 +15,7 @@ limitations under the License. */
 #pragma once
 
 #include <stddef.h>  // for size_t
+#include <unordered_map>
 
 namespace paddle {
 namespace memory {
@@ -44,6 +45,8 @@ class CPUAllocator : public SystemAllocator {
 #ifdef PADDLE_WITH_CUDA
 class GPUAllocator : public SystemAllocator {
  public:
+  std::unordered_map<int, uint64_t> s_memoryMap;
+
   explicit GPUAllocator(int gpu_id) : gpu_id_(gpu_id) {}
 
   virtual void* Alloc(size_t* index, size_t size);
-- 
GitLab