saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-09 17:28:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Not to calc memory for inference (#8935)

Browse source
This commit is contained in:
Wei-Sheng Chin 2021-09-02 09:49:54 -07:00 committed by GitHub
parent fbb6f0f599
commit f711d8992a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 11 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/core/framework/execution_frame.cc
View file

@ -404,6 +404,7 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const
ORT_TRY {
auto peak_size = mem_patterns_->patterns[i].PeakSize();
// Planning of one memory type should only happen once.
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
ORT_ENFORCE(
static_activation_memory_sizes_in_byte_.find(location.name) ==
static_activation_memory_sizes_in_byte_.end(),
@ -413,6 +414,7 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const
// static_activation_memory_in_bytes_ is max virtual memory size the planner computes.
// Memory dynamically allocated when executing kernels is not recorded using this field.
static_activation_memory_sizes_in_byte_[location.name] = peak_size;
#endif
buffer = alloc->Alloc(peak_size);
// handle allocator that doesn't throw
if (buffer == nullptr) {
@ -547,12 +549,12 @@ Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(OrtValue& ort_va
}
{
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
// This code block is not thread-safe.
// Dynamic activation size would be accessed by multiple threads
// if parallel executor is used.
std::unique_lock<std::mutex> lock(mtx_);
dynamic_activation_memory_sizes_in_byte_[location.name] += size;
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::SetDynamicAllocation(ort_value_index);
#endif
}

4
onnxruntime/core/framework/execution_frame.h
View file

@ -158,6 +158,7 @@ class ExecutionFrame final : public IExecutionFrame {
// If the retrival is sucessful, this function returns true and false otherwise.
bool TryGetInferredShape(int index, TensorShape& shape) const override;
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
// Return the size of virtual memory allocated in runtime.
// The memory is usually used for activations in forward and backward passes.
const std::unordered_map<std::string, size_t>& GetDynamicMemorySizeInfo() {
@ -177,6 +178,7 @@ class ExecutionFrame final : public IExecutionFrame {
// std::unique_lock<std::mutex> lock(mtx_);
return static_activation_memory_sizes_in_byte_;
}
#endif
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ExecutionFrame);
@ -225,6 +227,7 @@ class ExecutionFrame final : public IExecutionFrame {
// inferred_shapes_ is generated together with mem_patterns_.
std::unordered_map<int, TensorShape> inferred_shapes_;
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
// Size of virtual memory allocated before any kernel execution.
// This field is not physical memory size.
// static_activation_memory_sizes_in_byte_[location] is the static memory consumption on "location".
@ -239,5 +242,6 @@ class ExecutionFrame final : public IExecutionFrame {
// Mutex which should be acquired when executing non-thread-safe member functions.
// A current example is the tracker of dynamic memory allocation.
mutable std::mutex mtx_;
#endif
};
} // namespace onnxruntime

2
onnxruntime/core/framework/orttraining_partial_executor.cc
View file

@ -506,6 +506,7 @@ Status PartialExecutor::Execute(const SessionState& session_state, const std::ve
session_state.Profiler().EndTimeAndRecordEvent(profiling::SESSION_EVENT, "SequentialExecutor::Execute", tp);
}
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
for (auto i : frame.GetStaticMemorySizeInfo()) {
LOGS(logger, INFO) << "[Memory] ExecutionFrame statically allocates "
<< i.second << " bytes for " << i.first << std::endl;
@ -515,6 +516,7 @@ Status PartialExecutor::Execute(const SessionState& session_state, const std::ve
LOGS(logger, INFO) << "[Memory] ExecutionFrame dynamically allocates "
<< i.second << " bytes for " << i.first << std::endl;
}
#endif
return Status::OK();
}

2
onnxruntime/core/framework/sequential_executor.cc
View file

@ -484,6 +484,7 @@ Status SequentialExecutor::Execute(const SessionState& session_state, const std:
session_state.Profiler().EndTimeAndRecordEvent(profiling::SESSION_EVENT, "SequentialExecutor::Execute", tp);
}
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
for (auto i : frame.GetStaticMemorySizeInfo()) {
LOGS(logger, INFO) << "[Memory] ExecutionFrame statically allocates "
<< i.second << " bytes for " << i.first << std::endl;
@ -493,6 +494,7 @@ Status SequentialExecutor::Execute(const SessionState& session_state, const std:
LOGS(logger, INFO) << "[Memory] ExecutionFrame dynamically allocates "
<< i.second << " bytes for " << i.first << std::endl;
}
#endif
return Status::OK();
}