Memory Saturation and Kernel Restart when Converting to np.array in Python

Question

I'm encountering a memory issue when converting my processed data to a numpy array. I have 57GB of RAM, but the RAM saturates quickly and the kernel restarts at np.array(processed_X). Here is my code:

import numpy as np
import scipy.signal
from skimage.transform import resize
from tqdm import tqdm

def apply_stft(signal, nperseg=256, noverlap=128):
    f, t, Zxx = scipy.signal.stft(signal, nperseg=nperseg, noverlap=noverlap)
    return np.abs(Zxx)

def resize_stft(stft_result, output_shape=(224, 224)):
    return resize(stft_result, output_shape, mode='constant')

def process_dataset(X):
    processed_X = []
    for example in tqdm(X):
        stft_results = [apply_stft(example[:, i]) for i in range(example.shape[1])]
        resized_stft = [resize_stft(stft) for stft in stft_results]
        stacked_stft = np.stack(resized_stft, axis=-1)
        processed_X.append(stacked_stft)
        del stft_results, resized_stft, stacked_stft
    print(len(processed_X))
    print('FINISHED !!!')
    return np.array(processed_X)

When the code reaches np.array(processed_X), my RAM usage spikes rapidly, eventually consuming all available memory and causing the kernel to restart. How can I avoid this memory issue?

Thank you for your help!

Steps Taken:

• Tried deleting intermediate variables to free up memory.
• Verified that the issue occurs at the np.array(processed_X) line.