diff --git a/src/io/paddle_mobile.h b/src/io/paddle_mobile.h index 74c11471566c3db8a37ea2d62e0496e5d40cb3b7..0c4dfa9782458b03c7d5c4660e15e89937142dc7 100644 --- a/src/io/paddle_mobile.h +++ b/src/io/paddle_mobile.h @@ -47,6 +47,10 @@ class PaddleMobile { * */ bool Load(const std::string &model_path, const std::string ¶_path, bool optimize = 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