Train test split: примеры (PYTHON)

Примеры использования train_test_split

Разделение одного массива данных:

from sklearn.model_selection import train_test_split
import numpy as np

X = np.arange(10).reshape(5, 2)
X_train, X_test = train_test_split(X, test_size=0.4, random_state=42)
print("Обучающая выборка:")
print(X_train)
print("Тестовая выборка:")
print(X_test)

Обучающая выборка:
[[4 5]
 [0 1]
 [6 7]]
Тестовая выборка:
[[2 3]
 [8 9]]

Разделение двух массивов (признаки и метки):

X = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]
y = [0, 1, 0, 1, 0]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
print(f"X_train: {X_train}")
print(f"X_test: {X_test}")
print(f"y_train: {y_train}")
print(f"y_test: {y_test}")

X_train: [[9, 10], [3, 4], [1, 2]]
X_test: [[7, 8], [5, 6]]
y_train: [0, 1, 0]
y_test: [1, 0]

Стратифицированное разделение:

X = np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])
y = np.array([0, 0, 0, 1, 1, 1])
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, stratify=y, random_state=42)
print(f"Распределение в исходных данных: {np.bincount(y)}")
print(f"Распределение в обучающих данных: {np.bincount(y_train)}")
print(f"Распределение в тестовых данных: {np.bincount(y_test)}")

Распределение в исходных данных: [3 3]
Распределение в обучающих данных: [2 2]
Распределение в тестовых данных: [1 1]

Альтернативные функции в Python

Библиотека scikit-learn предлагает несколько альтернативных методов разделения данных:

StratifiedShuffleSplit - предоставляет более гибкий контроль над стратифицированным разделением, позволяя выполнять несколько итераций разделения с разными случайными состояниями.

TimeSeriesSplit - предназначен для временных рядов, где важно сохранять временной порядок. Разделение происходит последовательно без перемешивания.

KFold, StratifiedKFold - реализуют k-блочную кросс-валидацию. Разделяют данные на k блоков, каждый из которых по очереди используется в качестве тестового набора.

GroupShuffleSplit - обеспечивает разделение с учетом групп, гарантируя, что все элементы одной группы попадут либо в обучающую, либо в тестовую выборку.

Выбор функции зависит от специфики задачи. train_test_split оптимальна для простого случайного разделения, тогда как кросс-валидационные сплиттеры предпочтительнее при ограниченном объеме данных или необходимости более надежной оценки модели.

Аналоги функции в других языках программирования

R

В языке R функция createDataPartition из пакета caret выполняет схожую задачу:

library(caret)
data(iris)
set.seed(42)
trainIndex <- createDataPartition(iris$Species, p = 0.8, list = FALSE)
trainData <- iris[trainIndex, ]
testData <- iris[-trainIndex, ]

Java (Weka)

Библиотека Weka предоставляет класс RandomSplit для разделения данных:

import weka.core.Instances;
import java.io.FileReader;

Instances data = new Instances(new FileReader("data.arff"));
data.randomize(new Random(42));
int trainSize = (int) Math.round(data.numInstances() * 0.8);
int testSize = data.numInstances() - trainSize;
Instances train = new Instances(data, 0, trainSize);
Instances test = new Instances(data, trainSize, testSize);

JavaScript (TensorFlow.js)

TensorFlow.js предлагает метод tf.data.Dataset для работы с данными:

const features = tf.tensor2d([[1, 2], [3, 4], [5, 6], [7, 8]]);
const labels = tf.tensor1d([0, 1, 0, 1]);
const splitRatio = 0.75;
const trainSize = Math.floor(features.shape[0] * splitRatio);
const [trainFeatures, testFeatures] = tf.split(features, [trainSize, features.shape[0] - trainSize], 0);

PHP (PHP-ML)

Библиотека PHP-ML содержит функцию stratisfiedSplit:

use Phpml\ModelSelection\StratifiedRandomSplit;
use Phpml\Dataset\ArrayDataset;

$dataset = new ArrayDataset($features, $labels);
$split = new StratifiedRandomSplit($dataset, 0.3);
$train = $split->getTrain();
$test = $split->getTest();

SQL

В SQL данные можно разделить с помощью случайной выборки:

-- Создание обучающей выборки (70%)
CREATE TABLE train AS
SELECT * FROM dataset
WHERE random() < 0.7;

-- Создание тестовой выборки (30%)
CREATE TABLE test AS
SELECT * FROM dataset
WHERE NOT EXISTS (SELECT 1 FROM train WHERE train.id = dataset.id);

Типичные ошибки при использовании

Указание несуществующего размера выборки:

from sklearn.model_selection import train_test_split
import numpy as np

try:
    X = np.array([[1, 2], [3, 4], [5, 6]])
    X_train, X_test = train_test_split(X, test_size=2.0)  # Больше 1.0
    print("Разделение успешно")
except ValueError as e:
    print(f"Ошибка: {e}")

Ошибка: test_size=2.0 should be either an int and represent the absolute number of test samples, or a float and represent the proportion of the dataset to include in the test split

Некорректное использование параметра stratify:

try:
    X = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
    y = np.array([0, 1])  # Длина не совпадает с X
    X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=42)
except ValueError as e:
    print(f"Ошибка: {e}")

Ошибка: The 'stratify' parameter should be of length equal to the number of samples in X, but got 2 samples and 4 entries in X.

Попытка стратифицированного разделения при непрерывной целевой переменной:

try:
    X = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
    y = np.array([1.5, 2.3, 3.1, 4.7])  # Непрерывные значения
    X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=42)
except ValueError as e:
    print(f"Ошибка: {e}")

Ошибка: Supported target types are: ('binary', 'multiclass'). Got 'continuous' instead.

Забытый параметр shuffle при работе с временными рядами:

# Для временных рядов обычно отключают перемешивание
X = np.array([[1], [2], [3], [4], [5]])
y = np.array([1, 2, 3, 4, 5])
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
print(f"X_train: {X_train.flatten()}")
print(f"X_test: {X_test.flatten()}")

X_train: [1 2 3 4]
X_test: [5]

Изменения в последних версиях

В версии scikit-learn 1.4 появилось предупреждение при использовании train_size и test_size одновременно, если их сумма не равна 1.0. Ранее такое использование могло приводить к неожиданным результатам.

Начиная с версии 0.24, параметр shuffle по умолчанию установлен в True для всех сплиттеров, включая train_test_split, что повышает согласованность API.

В версии 0.20 функция была перемещена из модуля sklearn.cross_validation в sklearn.model_selection. Старое расположение оставалось доступным для обратной совместимости, но в новых версиях рекомендуется использовать обновленный импорт.

В версии 1.1 улучшена обработка sparse матриц и pandas DataFrame, обеспечивающая более эффективное использование памяти при работе с большими наборами данных.

Расширенные примеры использования

Разделение нескольких массивов одновременно:

from sklearn.model_selection import train_test_split
import numpy as np

X1 = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
X2 = np.array([[10, 20], [30, 40], [50, 60], [70, 80]])
y = np.array([0, 1, 0, 1])

X1_train, X1_test, X2_train, X2_test, y_train, y_test = train_test_split(
    X1, X2, y, test_size=0.25, random_state=42
)

print(f"X1_train:\n{X1_train}")
print(f"X2_train:\n{X2_train}")
print(f"y_train: {y_train}")

X1_train:
[[5 6]
 [1 2]
 [7 8]]
X2_train:
[[50 60]
 [10 20]
 [70 80]]
y_train: [0 0 1]

Использование с pandas DataFrame:

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.DataFrame({
    'feature1': [1, 2, 3, 4, 5, 6, 7, 8],
    'feature2': [10, 20, 30, 40, 50, 60, 70, 80],
    'target': [0, 1, 0, 1, 0, 1, 0, 1]
})

X = df[['feature1', 'feature2']]
y = df['target']

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=42
)

print(f"Размер X_train: {X_train.shape}")
print(f"Размер X_test: {X_test.shape}")
print(f"Распределение в y_train:\n{y_train.value_counts()}")

Размер X_train: (5, 2)
Размер X_test: (3, 2)
Распределение в y_train:
1    3
0    2
Name: target, dtype: int64

Контроль баланса классов при малой выборке:

import numpy as np
from sklearn.model_selection import train_test_split

# Создание несбалансированного набора данных
X = np.array([[i, i*2] for i in range(100)])
y = np.array([0]*90 + [1]*10)  # 90% класса 0, 10% класса 1

# Без стратификации
X_train1, X_test1, y_train1, y_test1 = train_test_split(X, y, test_size=0.2, random_state=42)

# Со стратификацией
X_train2, X_test2, y_train2, y_test2 = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)

print("Без стратификации:")
print(f"  Обучающая: {np.bincount(y_train1)}")
print(f"  Тестовая: {np.bincount(y_test1)}")
print("\nСо стратификацией:")
print(f"  Обучающая: {np.bincount(y_train2)}")
print(f"  Тестовая: {np.bincount(y_test2)}")

Без стратификации:
  Обучающая: [73  7]
  Тестовая: [17  3]

Со стратификацией:
  Обучающая: [72  8]
  Тестовая: [18  2]

Разделение с сохранением индексов для pandas:

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.DataFrame({
    'A': range(10),
    'B': range(10, 20)
}, index=[f'sample_{i}' for i in range(10)])

train_idx, test_idx = train_test_split(df.index, test_size=0.3, random_state=42)
train_df = df.loc[train_idx]
test_df = df.loc[test_idx]

print(f"Индексы обучающей выборки: {list(train_idx)}")
print(f"Индексы тестовой выборки: {list(test_idx)}")

Индексы обучающей выборки: ['sample_9', 'sample_1', 'sample_6', 'sample_7', 'sample_3', 'sample_0', 'sample_5']
Индексы тестовой выборки: ['sample_8', 'sample_4', 'sample_2']