mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
367106f591
Summary: (no major functionality change, purely cosmetic) Closes https://github.com/caffe2/caffe2/pull/1098 Reviewed By: akyrola Differential Revision: D5638664 Pulled By: Yangqing fbshipit-source-id: bbae0589ec1afe938a186ccfce9f6ff1a986a5db
42 lines
1.1 KiB
C++
42 lines
1.1 KiB
C++
#include "caffe2/core/context.h"
|
|
#include "caffe2/core/logging.h"
|
|
#include "caffe2/core/tensor.h"
|
|
#include "caffe2/core/typeid.h"
|
|
|
|
CAFFE2_DEFINE_bool(
|
|
caffe2_report_cpu_memory_usage,
|
|
false,
|
|
"If set, print out detailed memory usage");
|
|
|
|
namespace caffe2 {
|
|
|
|
static std::unique_ptr<CPUAllocator> g_cpu_allocator(new DefaultCPUAllocator());
|
|
CPUAllocator* GetCPUAllocator() {
|
|
return g_cpu_allocator.get();
|
|
}
|
|
|
|
void SetCPUAllocator(CPUAllocator* alloc) {
|
|
g_cpu_allocator.reset(alloc);
|
|
}
|
|
|
|
MemoryAllocationReporter CPUContext::reporter_;
|
|
|
|
void MemoryAllocationReporter::New(void* ptr, size_t nbytes) {
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
size_table_[ptr] = nbytes;
|
|
allocated_ += nbytes;
|
|
LOG(INFO) << "Caffe2 alloc " << nbytes << " bytes, total alloc " << allocated_
|
|
<< " bytes.";
|
|
}
|
|
|
|
void MemoryAllocationReporter::Delete(void* ptr) {
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
auto it = size_table_.find(ptr);
|
|
CHECK(it != size_table_.end());
|
|
allocated_ -= it->second;
|
|
LOG(INFO) << "Caffe2 deleted " << it->second << " bytes, total alloc "
|
|
<< allocated_ << " bytes.";
|
|
size_table_.erase(it);
|
|
}
|
|
|
|
} // namespace caffe2
|