Eden-Kramer-Lab / spectral_connectivity

Frequency domain estimation and functional and directed connectivity analysis tools for electrophysiological data
GNU General Public License v3.0
121 stars 43 forks source link

[JOSS] GPU computations #53

Closed EtienneCmb closed 1 year ago

EtienneCmb commented 1 year ago

I tested GPU computations but I'm not sure I did it correctly as I couldn't observe differences in computing time (see notebook here). Do you have an idea why?

Also, in the doc and in the Multitaper file, you define an environmental variable SPECTRAL_CONNECTIVITY_ENABLE_GPU but in the Connectivity file it's SPEC_CON_ENABLE_GPU.

xref : JOSS (https://github.com/openjournals/joss-reviews/issues/4840)

edeno commented 1 year ago

Thanks for finding my error with the environmental variable. I fixed it in the code and will be fixed in the next release.

As for GPU performance, there are two things to be aware of:

  1. There is some care needed in benchmarking because of asynchronous execution (see https://docs.cupy.dev/en/stable/user_guide/performance.html#benchmarking). So you need to use the following to benchmark:
    
    from cupyx.profiler import benchmark

benchmark(gpu_computations, args=(data,), kwargs=kw, n_repeat=10)


2. There is some trickiness with using the environmental variable because the code uses the environmental variable at import time. Therefore simply changing the environmental variable or re-importing won't work. If you had two notebooks running different instances, you would see the speedup. Another way to handle this is to have a `use_gpu`argument which would explicitly control if the GPU is used, which is easier for the user to understand but more complicated code-wise. I will think about using this formulation. In any case, I will note this in the documentation.

This is what I get (using an A100 GPU) with the following parameters:
```python
n_trials = 100
n_times = 5000
n_roi = 10
sf = 128.

data = np.random.rand(n_times, n_trials, n_roi)

kw = dict(sampling_frequency=sf, n_tapers=7)
cpu_computations    :    CPU:2597133.193 us   +/-2930.501 (min:2593316.667 / max:2601036.531) us     GPU-0:2597173.315 us   +/-2924.542 (min:2593349.609 / max:2601070.557) us

gpu_computations    :    CPU:143584.640 us   +/-8525.303 (min:130068.105 / max:156613.369) us     GPU-0:151254.628 us   +/-8514.071 (min:137786.362 / max:164320.251) us
EtienneCmb commented 1 year ago

Ok, thanks for the explanations and it's nice to see performance improvements by using GPU.

Actually, I had the same problem of changing environmental variables before/after importing the package. If you find a workaround, please let me know.

Best,