使用sklearn train_test_split处理3D数组时内核崩溃求助

Hey there, sorry to hear you're stuck with this kernel crash when using train_test_split on your large 3D dataset. Let's walk through why this is happening and the practical fixes you can try right away.

Why This Happens #

First, let's do a quick memory calculation for your data: your feature array (1860000, 144, 12) takes up ~128GB if stored as float32, or ~256GB as float64 (numpy's default for many cases). When train_test_split runs on this full array, it often creates internal copies or intermediate arrays that push your system's memory over the limit—leading to the kernel crashing without warning. Updating packages won't fix this core memory issue, so we need to adjust how we handle the data.

Practical Solutions #

1. Split Indices First, Then Extract Data (Most Efficient Fix)

Instead of passing the entire huge array to train_test_split, work with sample indices instead. This way, the function only processes a tiny 1D array of indices, and you can slice the original data without copying it (numpy creates memory views when slicing, not full copies):

import numpy as np
from sklearn.model_selection import train_test_split

# Generate indices for all samples
sample_indices = np.arange(X.shape[0])

# Split indices instead of the full arrays
train_idx, test_idx = train_test_split(
    sample_indices, 
    test_size=0.2,  # Adjust to your needs
    random_state=42,
    stratify=y  # Add this if you need stratified splits for classification
)

# Extract training/test data using the indices
X_train = X[train_idx]
X_test = X[test_idx]
y_train = y[train_idx]
y_test = y[test_idx]

This cuts down the memory load of train_test_split to almost nothing, and slicing the original arrays won't duplicate memory unless you modify the sliced arrays later.

2. Reduce Data Precision

If your use case allows it, downcasting your data to a smaller dtype can halve or quarter the memory footprint:

# Convert features to float32 (if you don't need float64 precision)
X = X.astype(np.float32)

# If labels are integers, use int32 instead of int64
y = y.astype(np.int32)

This simple change might free up enough memory to let train_test_split run normally without index splitting.

3. Use Memory-Mapped Arrays (For Extremely Large Data)

If even downcasting isn't enough (your data is way larger than available RAM), use numpy's memory-mapped arrays. These let you access data from disk without loading everything into memory at once:

# Save your array to a memory-mapped file
X_mmap = np.memmap(
    'features_mmap.npy',
    dtype=np.float32,
    mode='w+',
    shape=X.shape
)
X_mmap[:] = X[:]  # Write data to the mmap file
del X  # Free up original array memory
X_mmap.flush()  # Ensure data is written to disk

# Do the same for labels
y_mmap = np.memmap(
    'labels_mmap.npy',
    dtype=np.int32,
    mode='w+',
    shape=y.shape
)
y_mmap[:] = y[:]
del y
y_mmap.flush()

# Now split indices and slice the memory-mapped arrays
sample_indices = np.arange(X_mmap.shape[0])
train_idx, test_idx = train_test_split(sample_indices, test_size=0.2, random_state=42)

X_train = X_mmap[train_idx]
X_test = X_mmap[test_idx]
y_train = y_mmap[train_idx]
y_test = y_mmap[test_idx]

This works even if your data is several times larger than your available RAM, as numpy loads only the slices you need into memory on demand.

4. Check System Memory & Swap

If you're still crashing, verify your system's available memory (use htop or top on Linux/macOS, Task Manager on Windows). If RAM is fully exhausted, adding temporary swap space can act as a buffer, though this will slow down operations. This is a band-aid, but it might get you through the split if you need a quick fix.

Final Notes #

Double-check your scikit-learn version (print(sklearn.__version__)) to ensure you're on the latest stable release—older versions had rare memory leaks with large arrays, but this is less likely than the memory overload issue we've focused on.

内容的提问来源于stack exchange，提问作者MadPhil