Caffe2 - C++ API
A deep learning, cross platform ML framework
allocator.cc
1 
17 #include "caffe2/core/context.h"
18 #include "caffe2/core/logging.h"
19 #include "caffe2/core/tensor.h"
20 #include "caffe2/core/typeid.h"
21 
22 CAFFE2_DEFINE_bool(
23  caffe2_report_cpu_memory_usage,
24  false,
25  "If set, print out detailed memory usage");
26 
27 CAFFE2_DEFINE_bool(
28  caffe2_cpu_allocator_do_zero_fill,
29  true,
30  "If set, do memory zerofilling when allocating on CPU");
31 
32 namespace caffe2 {
33 
34 void NoDelete(void*) {}
35 
36 static std::unique_ptr<CPUAllocator> g_cpu_allocator(new DefaultCPUAllocator());
37 CPUAllocator* GetCPUAllocator() {
38  return g_cpu_allocator.get();
39 }
40 
41 void SetCPUAllocator(CPUAllocator* alloc) {
42  g_cpu_allocator.reset(alloc);
43 }
44 
45 MemoryAllocationReporter CPUContext::reporter_;
46 
47 void MemoryAllocationReporter::New(void* ptr, size_t nbytes) {
48  std::lock_guard<std::mutex> guard(mutex_);
49  size_table_[ptr] = nbytes;
50  allocated_ += nbytes;
51  LOG(INFO) << "Caffe2 alloc " << nbytes << " bytes, total alloc " << allocated_
52  << " bytes.";
53 }
54 
55 void MemoryAllocationReporter::Delete(void* ptr) {
56  std::lock_guard<std::mutex> guard(mutex_);
57  auto it = size_table_.find(ptr);
58  CHECK(it != size_table_.end());
59  allocated_ -= it->second;
60  LOG(INFO) << "Caffe2 deleted " << it->second << " bytes, total alloc "
61  << allocated_ << " bytes.";
62  size_table_.erase(it);
63 }
64 
65 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.