Skip to content

Commit dc5a7b9

Browse files
committed
fix default number of threads when inference with or without MKLDNN
test=develop
1 parent 5f2e837 commit dc5a7b9

File tree

2 files changed

+10
-0
lines changed

2 files changed

+10
-0
lines changed

paddle/fluid/inference/api/analysis_predictor.cc

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,11 @@
2525
#include "paddle/fluid/inference/api/paddle_inference_api.h"
2626
#include "paddle/fluid/inference/api/paddle_inference_pass.h"
2727
#include "paddle/fluid/inference/utils/singleton.h"
28+
#include "paddle/fluid/platform/cpu_helper.h"
2829
#include "paddle/fluid/platform/profiler.h"
2930

3031
DECLARE_bool(profile);
32+
DECLARE_int32(paddle_num_threads);
3133

3234
namespace paddle {
3335

@@ -47,6 +49,9 @@ bool AnalysisPredictor::Init(
4749
}
4850
#endif
4951

52+
// no matter with or without MKLDNN
53+
paddle::platform::SetNumThreads(FLAGS_paddle_num_threads);
54+
5055
if (config_.use_gpu) {
5156
place_ = paddle::platform::CUDAPlace(config_.device);
5257
LOG(WARNING) << "ir optimize only supports CPU currently, enable_ir_optim "

paddle/fluid/inference/api/api_impl.cc

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,11 @@ limitations under the License. */
2323
#include "paddle/fluid/framework/feed_fetch_method.h"
2424
#include "paddle/fluid/inference/api/api_impl.h"
2525
#include "paddle/fluid/inference/api/helper.h"
26+
#include "paddle/fluid/platform/cpu_helper.h"
2627
#include "paddle/fluid/platform/profiler.h"
2728

2829
DEFINE_bool(profile, false, "Turn on profiler for fluid");
30+
DECLARE_int32(paddle_num_threads);
2931

3032
namespace paddle {
3133
namespace {
@@ -72,6 +74,9 @@ bool NativePaddlePredictor::Init(
7274
}
7375
#endif
7476

77+
// no matter with or without MKLDNN
78+
paddle::platform::SetNumThreads(FLAGS_paddle_num_threads);
79+
7580
if (config_.use_gpu) {
7681
place_ = paddle::platform::CUDAPlace(config_.device);
7782
} else {

0 commit comments

Comments
 (0)