1 #include "caffe2/core/tensor.h" 3 #include "caffe2/core/blob_stats.h" 9 "If set, keeps memory when a tensor is shrinking its size.");
12 caffe2_max_keep_on_shrink_memory,
14 "The maximum memory in bytes to keep on shrink, if the difference between " 15 "tensor sizes is bigger than this then tensor will be reset.");
21 TensorPrinter::TensorPrinter(
22 const std::string& tensor_name,
23 const std::string& file_name,
25 : to_file_(!file_name.empty()),
26 limit_(limit ? limit : k_limit_default_),
27 tensor_name_(tensor_name) {
31 log_file_.reset(
new std::ofstream(
32 file_name, std::ofstream::out | std::ofstream::trunc));
35 "Failed to open TensorPrinter file ",
38 log_file_->rdstate());
42 TensorPrinter::~TensorPrinter() {
43 if (log_file_.get()) {
48 static CaffeMap<CaffeTypeId, TypeCall> type_call_registry_ {
49 {TypeMeta::Id<Tensor<CPUContext>>(), GetTensorType<CPUContext>}
52 TypeCall GetTypeCallFunction(CaffeTypeId
id) {
53 auto f = type_call_registry_.find(
id);
54 if (f == type_call_registry_.end()) {
60 void RegisterTypeCallFunction(CaffeTypeId
id, TypeCall c) {
61 type_call_registry_[id] = c;
64 static CaffeMap<CaffeTypeId, TensorInfoCall> tensor_info_call_registry_{
65 {TypeMeta::Id<Tensor<CPUContext>>(), GetTensorInfo<CPUContext>}};
67 TensorInfoCall GetTensorInfoFunction(CaffeTypeId
id) {
68 auto f = tensor_info_call_registry_.find(
id);
69 if (f == tensor_info_call_registry_.end()) {
75 void RegisterTensorInfoFunction(CaffeTypeId
id, TensorInfoCall c) {
76 tensor_info_call_registry_[id] = c;
81 struct TensorCPUStatGetter : BlobStatGetter {
82 size_t sizeBytes(
const Blob& blob)
const override {
83 const auto& tensor = blob.Get<TensorCPU>();
84 auto nbytes = tensor.nbytes();
85 if (nbytes > 0 && tensor.IsType<std::string>()) {
86 const auto* data = tensor.data<std::string>();
87 for (
size_t i = 0; i < tensor.size(); ++i) {
88 nbytes += data[i].size();
94 REGISTER_BLOB_STAT_GETTER(TensorCPU, TensorCPUStatGetter);
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Commandline flags support for Caffe2.