train_test_split is a scikit-learn function for partitioning data into training and test subsets. In an ML workflow, the split must also be usable downstream: features need their corresponding targets, and pipeline components need accessible dataset artifacts. section 7.3.1 section 4.1.2
Split encoded features and targets
After encoding categorical features and the target, split both together: `X_train, X_test, y_train, y_test = train_test_split(dummyfied_df, target, train_size=0.80, shuffle=True)`. The resulting datasets can be saved to object storage and tracked with the model run for later loading or debugging. section 4.1.2
For YOLO data, two subsets are not enough: training, validation, and test outputs are required. Apply train_test_split twice—first to separate training data, then to divide the remainder into validation and test data. The example targets a 75/15/10 split and passes `random_state` for reproducibility. Keep each image with its label, then copy the partitions into separate `images` and `labels` directories. section 7.3.1
Listing 7.19 Using multiple Output[Dataset] and inner imports
@dsl.component(
packages_to_install=["scikit-learn"],
base_image="python:3.11"
)
def split_dataset(
random_state: int,
input_dataset: Input[Dataset], #1
train_dataset: Output[Dataset], #2
validation_dataset: Output[Dataset], #2
test_dataset: Output[Dataset] #2
):
import os
import glob
import shutil
from sklearn.model_selection import train_test_split
The Kubeflow component installs scikit-learn, imports train_test_split inside the component, and declares separate Dataset outputs for all three partitions. section 7.3.1
Listing 7.21 Helper function to organize dataset splits
def split_dataset(...):
for dataset_output, x_files, y_files in [ #1
(train_dataset, x_train, y_train),
(validation_dataset, x_val, y_val),
(test_dataset, x_test, y_test)
]:
os.makedirs(
os.path.join(
dataset_output.path,
"images"
),
exist_ok=True
) #2
os.makedirs(
os.path.join(dataset_output.path, "labels"),
exist_ok=True
)
for src in x_files: #3
dest = os.path.join(
dataset_output.path,
"images",
os.path.basename(src)
)
shutil.copy2(src, dest)
for src in y_files: #4
dest = os.path.join(
dataset_output.path,
"labels",
os.path.basename(src)
)
shutil.copy2(src, dest)
The output helper creates the image and label directories and copies each partition’s files into its Dataset artifact. section 7.3.1
