基于Numpy实现带类别分布保持的5折分层抽样与划分

Hey there! Let's build a pure NumPy function that handles dataset loading, stratified 5-fold splitting, and train-test splits for each fold while keeping the class distribution consistent. I'll integrate the category distribution stats and basic k-fold logic you already have, and add the stratified partitioning part.

Complete Implementation #

import numpy as np

def stratified_5fold_process(dataset_path):
    # Step 1: Load dataset into NumPy array
    data = np.genfromtxt(dataset_path, delimiter=',', dtype=None, encoding='utf-8')
    # Separate features and labels (last column is assumed to be the class label)
    X = data[:, :-1]
    y = data[:, -1]
    
    # Step 2: Calculate & display original class distribution (your existing stats logic)
    unique_classes, class_counts = np.unique(y, return_counts=True)
    class_dist = dict(zip(unique_classes, class_counts))
    print("Original dataset class distribution:", class_dist)
    
    # Step 3: Create stratified 5 folds
    folds = [[] for _ in range(5)]
    for cls in unique_classes:
        # Get all indices belonging to the current class
        cls_indices = np.where(y == cls)[0]
        # Shuffle indices to avoid ordering bias
        np.random.shuffle(cls_indices)
        # Split the class's indices into 5 roughly equal parts
        cls_splits = np.array_split(cls_indices, 5)
        # Assign each part to the corresponding fold
        for fold_idx, split in enumerate(cls_splits):
            folds[fold_idx].extend(split)
    
    # Convert fold lists to NumPy arrays for easier indexing
    folds = [np.array(fold) for fold in folds]
    
    # Step 4: Split each fold into 80% train / 20% test (stratified)
    fold_train_test = []
    for fold in folds:
        # Extract labels for the current fold
        fold_y = y[fold]
        train_indices = []
        test_indices = []
        
        for cls in unique_classes:
            # Get indices of the current class within the fold
            cls_fold_indices = fold[np.where(fold_y == cls)[0]]
            # Calculate 80% split point
            split_point = int(len(cls_fold_indices) * 0.8)
            # Shuffle before splitting to ensure randomness
            np.random.shuffle(cls_fold_indices)
            # Assign to train and test sets
            train_indices.extend(cls_fold_indices[:split_point])
            test_indices.extend(cls_fold_indices[split_point:])
        
        # Package results for the fold
        fold_train_test.append({
            'train_indices': np.array(train_indices),
            'test_indices': np.array(test_indices),
            'train_data': X[train_indices],
            'train_labels': y[train_indices],
            'test_data': X[test_indices],
            'test_labels': y[test_indices]
        })
    
    return fold_train_test

# Example usage with your modified E. coli dataset
if __name__ == "__main__":
    fold_results = stratified_5fold_process("ecoli_modified.csv")
    
    # Verify stratification works by checking fold 1's train/test distribution
    print("\nClass distribution in Fold 1 Train set:")
    train_classes, train_counts = np.unique(fold_results[0]['train_labels'], return_counts=True)
    print(dict(zip(train_classes, train_counts)))
    
    print("\nClass distribution in Fold 1 Test set:")
    test_classes, test_counts = np.unique(fold_results[0]['test_labels'], return_counts=True)
    print(dict(zip(test_classes, test_counts)))

Key Details & Explanations #

• Dataset Loading: Uses np.genfromtxt() to load CSV data. Adjust delimiter if your dataset uses tabs or other separators.
• Stratified Fold Creation:
  • For each class, we shuffle its sample indices first to eliminate any ordering bias from the original dataset.
  • np.array_split() handles uneven class sizes gracefully, ensuring each fold gets a proportional number of samples from every class.
  • This guarantees each fold has nearly identical class distribution to the original dataset.
• Train-Test Split per Fold:
  • Repeats the stratified logic for individual folds, so the 80/20 split also preserves the fold's class ratio.
  • Returns a list of dictionaries, making it easy to access train/test features, labels, and indices for each fold.
• Validation: The example code prints class distributions for the first fold's train/test sets, so you can confirm stratification works as expected.

Quick Notes #

• Ensure your dataset has the class label as the last column (adjust X and y slicing if your format differs).
• Shuffling is added at multiple steps to ensure randomness without breaking class balance.
• For tiny datasets, you might see minor class count differences across folds, but this implementation minimizes that as much as possible.

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