diff --git a/src/io/paddle_mobile.h b/src/io/paddle_mobile.h index 31a920790264025bfa9aa996cc396edc9f8c3f24..f0f62fd3b39bc895e3cbfeeee21aadb9e753e5f0 100644 --- a/src/io/paddle_mobile.h +++ b/src/io/paddle_mobile.h @@ -48,6 +48,9 @@ class PaddleMobile { bool Load(const std::string &model_path, const std::string ¶_path, bool optimize = false, bool quantification = false, int batch_size = 1); + /* + * @b 设置线程数, 当 cmake 中开启 openmp 时生效 + * */ void SetThreadNum(int num); /* diff --git a/src/operators/kernel/central-arm-func/pool_arm_func.h b/src/operators/kernel/central-arm-func/pool_arm_func.h index 12aa01507d83c5051f4b462fc4607d71a707f06d..6179df5b0c11ad2a2e19384989029696e9d6c266 100644 --- a/src/operators/kernel/central-arm-func/pool_arm_func.h +++ b/src/operators/kernel/central-arm-func/pool_arm_func.h @@ -77,11 +77,15 @@ void PoolCompute(const PoolParam ¶m) { } else if (ksize[0] == 2 && ksize[0] == ksize[1]) { #if __ARM_NEON +#if __aarch64__ + PoolBasic(pooling_type, ksize, strides, paddings, in_x, out); +#else if (pooling_type == "max") { math::Pool2x2Max(strides, paddings, in_x, out); } else if (pooling_type == "avg") { math::Pool2x2Avg(strides, paddings, in_x, out); } +#endif #else PoolBasic(pooling_type, ksize, strides, paddings, in_x, out); #endif // __ARM_NEON