Scikit Learn

scikit-learn Best Practices

This guide outlines our team's definitive best practices for using and extending scikit-learn. Adhering to these rules ensures consistent, reproducible, and production-ready machine learning code.

1. Code Organization and Structure

1.1. Always Use Pipelines for Preprocessing and Models

Pipelines are mandatory. They prevent data leakage, ensure consistent transformations across training and inference, and simplify hyperparameter tuning.

❌ BAD: Inconsistent manual transformations

from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import numpy as np

X, y = np.random.rand(100, 5), np.random.rand(100)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
model = LinearRegression().fit(X_train_scaled, y_train)

# Forgetting to scale X_test leads to incorrect predictions
y_pred = model.predict(X_test)
print(f"MSE (BAD): {mean_squared_error(y_test, y_pred):.2f}")

✅ GOOD: Encapsulate all steps in a Pipeline

from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import numpy as np

X, y = np.random.rand(100, 5), np.random.rand(100)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

# All steps are chained and applied consistently
model = make_pipeline(StandardScaler(), LinearRegression())
model.fit(X_train, y_train)
y_pred = model.predict(X_test)
print(f"MSE (GOOD): {mean_squared_error(y_test, y_pred):.2f}")

1.2. Custom Estimators Must Adhere to the scikit-learn API

When creating custom transformers or models, strictly follow the scikit-learn estimator API for seamless integration with pipelines and model selection tools.

• Inherit from BaseEstimator and relevant mixins (TransformerMixin, ClassifierMixin, RegressorMixin).
• __init__ stores only hyperparameters as keyword arguments with defaults. No logic or data processing here.
• fit(self, X, y=None) learns parameters from data. Store learned attributes with a trailing underscore (e.g., self.mean_). Return self.
• transform(self, X) (for transformers) or predict(self, X)/predict_proba(self, X) (for predictors) applies the learned transformation/prediction.
• Implement get_params and set_params (usually inherited from BaseEstimator).

❌ BAD: Non-compliant custom estimator

class BadCustomScaler:
    def __init__(self, scale_factor):
        # Logic in __init__
        if scale_factor <= 0:
            raise ValueError("Scale factor must be positive")
        self.scale_factor = scale_factor
        self.mean_ = None # Should be learned in fit

    def fit(self, X):
        self.mean_ = X.mean(axis=0)
        return self

    def transform(self, X):
        if self.mean_ is None:
            raise RuntimeError("Fit not called yet.")
        return (X - self.mean_) * self.scale_factor

# This won't work with GridSearchCV or Pipelines reliably

✅ GOOD: API-compliant custom estimator

from sklearn.base import BaseEstimator, TransformerMixin
import numpy as np

class GoodCustomScaler(BaseEstimator, TransformerMixin):
    def __init__(self, scale_factor: float = 1.0):
        # Only store hyperparameters, no logic
        self.scale_factor = scale_factor

    def fit(self, X: np.ndarray, y=None):
        # Input validation (optional, but good practice)
        X = self._validate_data(X)
        if self.scale_factor <= 0:
            raise ValueError("Scale factor must be positive")
        self.mean_ = X.mean(axis=0) # Learned attribute with underscore
        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        X = self._validate_data(X) # Ensure consistency
        if not hasattr(self, "mean_"):
            raise RuntimeError("Estimator not fitted. Call fit() first.")
        return (X - self.mean_) * self.scale_factor

# This estimator integrates perfectly with scikit-learn utilities.

2. Common Patterns and Anti-patterns

2.1. Prevent Data Leakage: Split First, Fit on Train Only

Always split your data into training and test sets before any preprocessing. fit and fit_transform methods must only be called on the training data. transform can then be called on both training and test data.

❌ BAD: Data leakage by fitting on all data

from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import numpy as np

X, y = np.random.rand(100, 5), np.random.rand(100)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

scaler = StandardScaler()
# Fitting on all data (X) leaks information from the test set
X_scaled = scaler.fit_transform(X)
X_train_scaled = X_scaled[:len(X_train)]
X_test_scaled = X_scaled[len(X_train):]

✅ GOOD: Correct data splitting and transformation

from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import numpy as np

X, y = np.random.rand(100, 5), np.random.rand(100)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

scaler = StandardScaler()
# Fit only on training data
X_train_scaled = scaler.fit_transform(X_train)
# Transform both train and test using parameters learned from train
X_test_scaled = scaler.transform(X_test)

2.2. Reproducibility: Always Set random_state

For any estimator or utility that involves randomness (e.g., train_test_split, RandomForestClassifier, KMeans), explicitly set the random_state parameter for reproducible results.

❌ BAD: Non-reproducible results

from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
# Results will change on each run
X_train, X_test, y_train, y_test = train_test_split(X, y)
model = RandomForestClassifier()

✅ GOOD: Reproducible results

from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier

# Use a fixed integer for random_state
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)
model = RandomForestClassifier(random_state=42)

3. Performance Considerations

3.1. Leverage Cython for Performance-Critical Components

When extending scikit-learn with custom, performance-critical algorithms, use Cython and follow scikit-learn's internal conventions.

• Disable bounds checking and wraparound for production code.
• Use sklearn.utils._typedefs for explicit type declarations.
• Use sklearn.utils._openmp_helpers for OpenMP routines.
• Prefer memoryviews over cnp.ndarray when possible.

# Example of Cython directives and imports (in a .pyx file)
# distutils: language=c
# distutils: extra_compile_args=-fopenmp
# distutils: extra_link_args=-fopenmp
# cython: boundscheck=False
# cython: wraparound=False
# cython: cdivision=True

from cython.parallel import prange
from sklearn.utils._typedefs cimport float64
from sklearn.utils._openmp_helpers cimport _get_num_threads

cdef void my_fast_function(float64[:] data, int n_threads) nogil:
    cdef int i
    with nogil:
        for i in prange(data.shape[0], num_threads=n_threads, schedule='static'):
            data[i] *= 2.0

4. Common Pitfalls and Gotchas

4.1. Avoid Overfitting on Validation Data During Hyperparameter Tuning

Use GridSearchCV or RandomizedSearchCV with a pipeline to ensure that cross-validation folds are correctly handled and preprocessing steps are refitted for each fold, preventing data leakage.

❌ BAD: Tuning hyperparameters outside a pipeline

from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, train_test_split
import numpy as np

X, y = np.random.rand(100, 5), np.random.randint(0, 2, 100)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

# Scaling outside of GridSearchCV means the scaler sees all training data
# before cross-validation, leading to optimistic scores.
scaler = StandardScaler().fit(X_train)
X_train_scaled = scaler.transform(X_train)

param_grid = {'C': [0.1, 1.0, 10.0]}
grid_search = GridSearchCV(LogisticRegression(random_state=42), param_grid, cv=3)
grid_search.fit(X_train_scaled, y_train)
# Scores here are optimistically biased

✅ GOOD: Tuning hyperparameters with a pipeline

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, train_test_split
import numpy as np

X, y = np.random.rand(100, 5), np.random.randint(0, 2, 100)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

# Define the pipeline including preprocessing and the model
pipeline = Pipeline([
    ('scaler', StandardScaler()),
    ('logreg', LogisticRegression(random_state=42))
])

param_grid = {'logreg__C': [0.1, 1.0, 10.0]} # Note the 'estimator__param' syntax
grid_search = GridSearchCV(pipeline, param_grid, cv=3)
grid_search.fit(X_train, y_train) # Fit the entire pipeline on X_train
# Scores are realistic as preprocessing is done correctly within each fold

5. Type Hints

5.1. Use Type Hints Extensively

All function signatures, class attributes, and complex variable assignments must include type hints. This improves code readability, enables static analysis, and reduces bugs.

❌ BAD: Untyped function

def calculate_metric(y_true, y_pred):
    # What are y_true and y_pred? Lists, arrays, Series?
    return (y_true == y_pred).mean()

✅ GOOD: Clearly typed function

import numpy as np
from typing import Union

def calculate_metric(y_true: Union[np.ndarray, list], y_pred: Union[np.ndarray, list]) -> float:
    """Calculates accuracy for binary classification."""
    y_true_arr = np.asarray(y_true)
    y_pred_arr = np.asarray(y_pred)
    return (y_true_arr == y_pred_arr).mean()

6. Virtual Environments

6.1. Isolate Project Dependencies with Virtual Environments

Always use a dedicated virtual environment (venv, conda, poetry, etc.) for each project. This prevents dependency conflicts and ensures consistent environments.

❌ BAD: Global pip install

pip install scikit-learn pandas numpy # Pollutes global environment

✅ GOOD: Project-specific virtual environment

# Using venv
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

# Using conda
conda create -n my_ml_env python=3.9 scikit-learn pandas numpy
conda activate my_ml_env

7. Packaging

7.1. Use scikit-learn-contrib Template for Shareable Estimators

If you develop a custom estimator intended for public use or sharing across projects, leverage the scikit-learn-contrib template. It provides a robust structure, testing, and CI/CD setup, ensuring your estimator is fully compliant and easily consumable.

❌ BAD: Ad-hoc custom estimator distribution

# Just a .py file with a custom estimator, no proper packaging
# Difficult for others to install, test, or integrate

✅ GOOD: Structured project using scikit-learn-contrib

# Follow the template to create a new project:
# https://github.com/scikit-learn-contrib/project-template
# This provides setup.py, tests, documentation structure, etc.

8. Testing Approaches

8.1. Implement Comprehensive Unit and Integration Tests

For custom estimators, write unit tests to verify individual methods and use sklearn.utils.estimator_checks.check_estimator to ensure API compliance. For pipelines, write integration tests that cover the entire workflow.

❌ BAD: No tests or only manual verification

# No tests for MyCustomEstimator, relying on manual runs
# to confirm it works.

✅ GOOD: Automated testing with check_estimator

import pytest
from sklearn.utils.estimator_checks import check_estimator
from sklearn.base import BaseEstimator, TransformerMixin
import numpy as np

class MyCustomTransformer(BaseEstimator, TransformerMixin):
    def __init__(self, offset: float = 0.0):
        self.offset = offset
    def fit(self, X, y=None):
        self._validate_data(X)
        return self
    def transform(self, X):
        self._validate_data(X, reset=False)
        return X + self.offset

# Test API compliance
@pytest.mark.parametrize(
    "estimator", [MyCustomTransformer()]
)
def test_all_estimators(estimator):
    return check_estimator(estimator)

# Add specific unit tests for logic
def test_my_custom_transformer_offset():
    X = np.array([[1, 2], [3, 4]])
    transformer = MyCustomTransformer(offset=10)
    transformer.fit(X)
    transformed_X = transformer.transform(X)
    assert np.array_equal(transformed_X, np.array([[11, 12], [13, 14]]))