Skip to content

Commit 056dd40

Browse files
committed
add initial memory flag in MB for infer
1 parent 431491a commit 056dd40

File tree

1 file changed

+8
-1
lines changed

1 file changed

+8
-1
lines changed

paddle/fluid/platform/cpu_info.cc

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,17 @@ limitations under the License. */
2121
#include <unistd.h>
2222
#endif
2323

24+
#include <algorithm>
2425
#include "gflags/gflags.h"
2526

2627
DEFINE_double(fraction_of_cpu_memory_to_use, 1,
2728
"Default use 100% of CPU memory for PaddlePaddle,"
2829
"reserve the rest for page tables, etc");
2930

31+
DEFINE_uint64(
32+
initial_cpu_memory_in_mb, 500,
33+
"Default initial 500MB of CPU memory for PaddlePaddle, in MD unit.");
34+
3035
DEFINE_double(
3136
fraction_of_cuda_pinned_memory_to_use, 0.5,
3237
"Default use 50% of CPU memory as the pinned_memory for PaddlePaddle,"
@@ -54,7 +59,9 @@ inline size_t CpuTotalPhysicalMemory() {
5459
size_t CpuMaxAllocSize() {
5560
// For distributed systems, it requires configuring and limiting
5661
// the fraction of memory to use.
57-
return FLAGS_fraction_of_cpu_memory_to_use * CpuTotalPhysicalMemory();
62+
return std::min(static_cast<size_t>(FLAGS_fraction_of_cpu_memory_to_use *
63+
CpuTotalPhysicalMemory()),
64+
FLAGS_initial_cpu_memory_in_mb * 1 << 20);
5865
}
5966

6067
size_t CpuMinChunkSize() {

0 commit comments

Comments
 (0)