Spike detection phase is very slow

[tayheau]

Hi everyone,
As the title says, im experimenting a slow spike detection phase using kilosort4 (~1.07s/it).
Im using a 24 hearth cpu, rtx4000ada and the whole thing on Windows-NT. The cpu is almost fully used but the GPU is not use except for it's memory which is fully used.
For the job's kwargs i just tweaked a bit the chunk_duration but that didnt really improved anything, and my n_jobs is set to 20.
torch gpu with the best cuda version for the gpu are installed.
My guess is that this phase was that this phase was only cpu based but even with different n_jobs or chunk_duration(so bigger IO flow) the s/it remained the same. I just think im missing something but i cant tell what (or maybe im just totally desillusionnal :D).

Thanks in advance for your help !

[chrishalcrow]

Hi @tayheau , just to clarify: you're running kilosort4 from spikeinterface, and the spike detection phase is taking a long time? Could you share a snippet of the code you're using?

How big is your recording? In terms of number of channels, and in terms of length?

[tayheau]

Hey @chrishalcrow thks for the fast response ! Yes of course so im on a 384 channels 30kHz 50 min long NP2.0 recording stored as a open ephys binary.
Actually i just wanted to "benchmark" my current workstation so the process is not very streamlined, here is the params of the last run :

{
   "sorter_name": "kilosort4",
   "sorter_params": {
       "batch_size": 60000,
       "nblocks": 1,
       "Th_universal": 9,
       "Th_learned": 8,
       "nt": 61,
       "shift": null,
       "scale": null,
       "batch_downsampling": 1,
       "artifact_threshold": Infinity,
       "nskip": 25,
       "whitening_range": 32,
       "highpass_cutoff": 300,
       "binning_depth": 5,
       "sig_interp": 20,
       "drift_smoothing": [
           0.5,
           0.5,
           0.5
       ],
       "nt0min": null,
       "dmin": null,
       "dminx": 32,
       "min_template_size": 10,
       "template_sizes": 5,
       "nearest_chans": 10,
       "nearest_templates": 100,
       "max_channel_distance": 32,
       "max_peels": 100,
       "templates_from_data": true,
       "n_templates": 6,
       "n_pcs": 6,
       "Th_single_ch": 6,
       "acg_threshold": 0.2,
       "ccg_threshold": 0.25,
       "cluster_neighbors": 10,
       "cluster_downsampling": 1,
       "max_cluster_subset": 25000,
       "x_centers": null,
       "cluster_init_seed": 5,
       "duplicate_spike_ms": 0.25,
       "position_limit": 100,
       "do_CAR": true,
       "invert_sign": false,
       "save_extra_vars": false,
       "save_preprocessed_copy": false,
       "torch_device": "auto",
       "bad_channels": null,
       "clear_cache": false,
       "do_correction": true,
       "skip_kilosort_preprocessing": false,
       "keep_good_only": false,
       "use_binary_file": true,
       "delete_recording_dat": true,
       "pool_engine": "process",
       "n_jobs": 20,
       "chunk_duration": "10s",
       "progress_bar": true,
       "mp_context": null,
       "max_threads_per_worker": 1
   }
}

Also, ive moved the OE folder to a local SSD just to be sure it's not due to I/O rate issue due to the connection.

Thanks again for your help and fast reponse !

[chrishalcrow]

Great, thanks - could you also share the code you use to actually run the sorting. Ideally from when you read in the recording to when you use run_sorter. Thanks!

[tayheau]

probeA_AP = se.read_openephys(oe_folder, stream_id='1')
protocol = {
    'preprocessing': {
        'bandpass_filter': {},
        'common_reference': {'operator': 'average'},
        'detect_and_remove_bad_channels': {},
    },
    'sorting': {
        'sorter_name': 'kilosort4',
        'folder': pathlib.Path(".") / "results_ks4",
        'remove_existing_folder': True,
        'verbose': True,
        'n_jobs': 20,
        'skip_kilosort_preprocessing': True,
    }
}
preprocessed = spre.apply_preprocessing_pipeline(probeA_AP, protocol["preprocessing"])
sorting = ss.run_sorter(recording=preprocessed, **protocol["sorting"])

[chrishalcrow]

Hey @tayheau sorry I didn't get back to you. Nothing immediately stuck on when I read your code. Did you manage to resolve it?

[tayheau]

Hey @chrishalcrow ! no problem man ! as far as i know, i just think this is a regular speed given the hardware and the data.

[chrishalcrow]

Cool cool - I'd expect a 1 hour NP recording to sort in about an hour (or two) with a ok GPU, and about 8 hours on a modern-ish CPU. We actually use CPUs in our lab, cause they're very available and cheap!

[tayheau]

ah ok so there definitly an issue with my setup: 40min NP1 recording was sorted in ~4h on a rtx4000ada. at first i though it could be due to gpu concurency since i use n>1 processes but i highly doubt that given the code quality ;) my second guess is maybe the env itself but every driver is up to date, the CUDA version i use is the recommanded one by nvidia and the graphic driver. I will investigate that.

[chrishalcrow]

Hmmm, it could be that use_binary_file is not working. This saves a copy of the preprocessed recording, so that kilosort doesn't have to preprocess on the fly (if it did, it would have to do this a lot, which can slow things down). I think if your recording is not binary compatible this fails silently. We should probably at least throw a warning... Ways to check if this is working: is there a spikeinterface tqdm progress bar before kilosort kicks in? Is there a recording.dat file created in sorter_output during sorting (note: this gets deleted once the sorting is completed)?>

[tayheau]

Yep, there is a recording.dat created during the sorting process in sorter_output. According to nvidia-smi, the GPU is used at 100% during sorting. The kilosort.spikedetect is the one who takes time (~1.10s/ it). Well i guess it's just definitly the expected perf

[chrishalcrow]

Just spoke to a colleague who had a problem with a slow deeplabcut. It turned out he had set his some low level parallelization off. This affected some parts of torch, which would affect both deeplabcut and kilosort. So maybe check some of that. This is what some AI bot told me to try.

import os

# List of environment variables to check
parallel_vars = [
    "MKL_NUM_THREADS",
    "OPENBLAS_NUM_THREADS",
    "OMP_NUM_THREADS",
    "BLOSC_NTHREADS",
    "BLOSC_MIN_DATA_SIZE",
    "NUMBA_NUM_THREADS",
]

for var_name in parallel_vars:
    # Use os.environ.get() to safely retrieve the variable's value
    # It returns None if the variable is not set, preventing errors.
    value = os.environ.get(var_name)

    if value is not None:
        print(f"**{var_name}**: Set to **{value}**")
    else:
        print(f"**{var_name}**: Not explicitly set in the environment.")

These things should not be 1!

[chrishalcrow]

I'd also recommend running kilosort directly with kilosort just to make sure spikeinterface isn't getting in the way and slowing something down (although I'm fairly confident it won't be!)