ARM-software / ComputeLibrary

The Compute Library is a set of computer vision and machine learning functions optimised for both Arm CPUs and GPUs using SIMD technologies.
2.87k stars 783 forks source link

How to Set the the num_threads when calling the runtime function? #1025

Closed GGGGxxxxxxxxr closed 1 year ago

GGGGxxxxxxxxr commented 1 year ago

H! I am building a WDSR model based on Arm Compute Library. As there are some operators (such as SkipLink) not supported by ArmComputeLibrary currently, I just manually call the runtime functions instead of using the GraphAPI.

I wonder how to set the num_threads during layer.run() when I don't utilize the GraphAPI?

Here is an example part of the code:

NEConvolutionLayer conv3_1;
conv3_1.configure(&layer2_out, &conv1_weight, &conv1_bias, &layer3_out, PadStrideInfo(1, 1, 1, 1));

conv3_1.run();

For example, I would like to use NUM_THREADS=1 in this conv3_1.run(), how can I set it?

Thanks!

GGGGxxxxxxxxr commented 1 year ago

Some update, I somehow managed to manually set the num_threads in CPPScheduler.cpp, but I am wondering whether I could have a function call to set it instead of modifying the backend source code too much.

morgolock commented 1 year ago

Hi @GGGGxxxxxxxxr

You can do this at runtime with the code below:

Scheduler::get().set_num_threads(threads_num);

For more details see: https://github.com/ARM-software/ComputeLibrary/blob/main/arm_compute/runtime/IScheduler.h#L146

Hope this helps