Skip to content

Shift subsets of channels at a time in phase_shift to reduce peak memory - #4904

Open
h-mayorquin wants to merge 1 commit into
SpikeInterface:mainfrom
h-mayorquin:phase_shift_memory
Open

h-mayorquin wants to merge 1 commit into
SpikeInterface:mainfrom
h-mayorquin:phase_shift_memory

Conversation

@h-mayorquin

Copy link
Copy Markdown
Contributor

This PR continues #4840 (Filter subsets of channels at a time to reduce peak memory) and #4903 (Filter subsets of channels at a time in gaussian_filter to reduce peak memory), now for phase_shift. It is the step that sets the peak memory of the usual Neuropixels chain, see the overview in #4837 (Process channels in banks in preprocessors to reduce RAM usage).

PhaseShiftRecordingSegment.get_traces casts the whole chunk to float64 and takes its FFT across all the channels at once. So an int16 chunk peaks at about 17 times its size. This PR shifts 16 channels at a time (reusing NUM_CHANNELS_PER_BLOCK from filter.py) and writes each block into a preallocated output. The peak goes down to 2 times the chunk for the phase shift and from 18 to 4 times for the whole chain. The output is identical to main for int16, float32 and float64 traces, for int16 traces with dtype="float32", at the edges of the recording and on a subset of channels.

With 1 s chunks it is also about 10% faster. With 10 s chunks it is about 25% slower, and about half of that comes from computing the rotations once per block instead of once per chunk.

peak memory (% of chunk) time
phase_shift, 1 s, main 1756% 0.30 s
phase_shift, 1 s, this PR 203% 0.26 s
phase_shift, 10 s, main 1644% 3.45 s
phase_shift, 10 s, this PR 197% 4.36 s
chain, 1 s, main 1824% 0.63 s
chain, 1 s, this PR 401% 0.60 s
chain, 10 s, main 1645% 6.55 s
chain, 10 s, this PR 401% 7.57 s

Neuropixels 2.0 recording, 384 channels at 30 kHz, int16. The chain is phase_shift, bandpass_filter to float32 and common_reference with the global median. The peak is from tracemalloc and includes the output array, the time is the best of 3 runs.

I am adding this now because I am doing some benchmarks in #4833 (Add num_chunk_per_batch to job kwargs and use it for zarr v3 shard writes), but it can come after #4563 (FIR phase-shift + extract apply_raised_cosine_taper from get_chunk_with_margin) and I am fine with fixing the conflicts.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance Performance issues/improvements preprocessing Related to preprocessing module

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants