Machine Learning on Big Data

Training models on massive datasets with Spark MLlib and beyond

When Your Data Won't Fit in RAM

Training ML models on 1GB of data? Scikit-learn works perfectly on your laptop. But what happens when you have 1TB of customer transactions, 10TB of sensor readings, or billions of user interactions? You need distributed machine learning. This lesson covers when and how to scale ML training from single-machine libraries (Scikit-learn) to distributed frameworks (Spark MLlib), and when to bring in deep learning powerhouses (TensorFlow, PyTorch) for massive neural networks.

The Scale Problem: Data vs Model Complexity

Different ML problems have different bottlenecks. Understanding this drives your choice of framework.

📊 Big Data, Simple Model

Problem: Billions of rows, basic algorithms

Example: Logistic regression on 1TB clickstream
Bottleneck: Data size (can't fit in RAM)
Model: Linear models, tree-based
Training Time: Hours to days

Use: Spark MLlib, H2O

🧠 Small Data, Complex Model

Problem: Millions of parameters, deep networks

Example: ResNet-50 on ImageNet (1M images)
Bottleneck: Model complexity (GPU compute)
Model: Deep neural networks
Training Time: Days on GPUs

Use: TensorFlow, PyTorch

💡 Key Decision: If your data fits in RAM (even if it takes 128GB), use single-node tools (Scikit-learn, PyTorch). Only use distributed frameworks (Spark MLlib) when data truly won't fit on the biggest available machine.

ML Framework Landscape

Framework	Best For	Data Size	Ease of Use	Speed
Scikit-learn	Traditional ML (trees, linear)	MB to GB (<100GB)	⭐⭐⭐⭐⭐	Fast (single-core)
Spark MLlib	Big data ML (distributed)	GB to TB (100GB+)	⭐⭐⭐	Moderate (parallelized)
TensorFlow	Production deep learning	GB to TB (w/ data pipeline)	⭐⭐⭐	Very fast (GPU)
PyTorch	Research, NLP, vision	GB to TB (w/ data pipeline)	⭐⭐⭐⭐	Very fast (GPU)
Keras	Fast prototyping (neural nets)	GB (TF/PyTorch backend)	⭐⭐⭐⭐⭐	Fast (GPU)

Scikit-learn: The Gold Standard for Traditional ML

Simple, consistent API for classical algorithms

Scikit-learn is the go-to library for traditional machine learning (regression, classification, clustering). It has the cleanest API, excellent documentation, and integrates perfectly with NumPy/Pandas. If your data fits in RAM, start here.

Key Strengths

🎯 Consistent API

fit(), predict(), score(), every model works the same way

📚 Rich Algorithms

100+ algorithms: SVM, Random Forest, XGBoost integration

🔧 Built-in Tools

Cross-validation, grid search, preprocessing all included

⚡ Production Ready

Mature, stable, well-tested, powers thousands of applications

Example: Classification with Random Forest

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report

# Load data (10M rows, fits in RAM)
df = pd.read_csv('customer_churn.csv')

# Prepare features and target
X = df[['age', 'tenure', 'monthly_charges', 'total_charges']]
y = df['churned']

# Split data
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Train model (simple!)
model = RandomForestClassifier(
    n_estimators=100,
    max_depth=10,
    random_state=42,
    n_jobs=-1  # Use all CPU cores
)
model.fit(X_train, y_train)

# Evaluate
y_pred = model.predict(X_test)
print(f"Accuracy: {accuracy_score(y_test, y_pred):.3f}")
print(classification_report(y_test, y_pred))

Result:
Training Time: ~3 minutes on 8 cores
Memory Usage: ~4GB RAM

Example: Hyperparameter Tuning

from sklearn.model_selection import GridSearchCV

# Define parameter grid
param_grid = {
    'n_estimators': [50, 100, 200],
    'max_depth': [5, 10, 15],
    'min_samples_split': [2, 5, 10]
}

# Grid search with cross-validation
grid_search = GridSearchCV(
    RandomForestClassifier(random_state=42),
    param_grid,
    cv=5,  # 5-fold cross-validation
    scoring='f1',
    n_jobs=-1,
    verbose=1
)

grid_search.fit(X_train, y_train)

print(f"Best params: {grid_search.best_params_}")
print(f"Best CV score: {grid_search.best_score_:.3f}")

Result:

Best params: {'max_depth': 10, 'min_samples_split': 5, 'n_estimators': 100}
Best CV score: 0.842

Note: Tested 27 combinations (3×3×3) with 5-fold CV = 135 models trained

When to Use Scikit-learn

Data fits in RAM (< 100GB on your machine)
Traditional ML algorithms (not deep learning)
Rapid prototyping and experimentation
Interpretable models (linear, trees)
Production deployments (stable, reliable)
When you need comprehensive preprocessing tools

Pros & Cons

✅ Pros

Easiest to learn and use
Excellent documentation
Consistent API across all models
Built-in preprocessing & validation
Production-ready (mature)
Great for structured/tabular data

❌ Cons

Single-machine only (not distributed)
Limited to ~100GB data
No GPU acceleration
No deep learning support
Can be slow for huge datasets
Not ideal for images/text/audio

Spark MLlib: Distributed Machine Learning

Scale to terabytes with distributed training

When your data exceeds 100GB and won't fit in RAM, Spark MLlib distributes training across a cluster. It provides familiar algorithms (linear models, trees, clustering) but operates on DataFrames that can span hundreds of machines.

Architecture: Distributed Training

┌──────────────────────────────────────────────────────┐
│          SPARK MLLIB DISTRIBUTED TRAINING            │
├──────────────────────────────────────────────────────┤
│                                                      │
│  Driver (Coordinator):                               │
│  ┌────────────────────────────────────┐              │
│  │  - Initialize model parameters     │              │
│  │  - Broadcast parameters to workers │              │
│  │  - Aggregate gradients             │              │
│  │  - Update model                    │              │
│  └────────┬───────────────────────────┘              │
│           │                                          │
│           │ Distributes data & params                │
│           ▼                                          │
│  ┌────────┴───────┬────────────┬────────┐            │
│  │                │            │        │            │
│  ▼                ▼            ▼        │            │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐  │            │
│ │Executor 1│ │Executor 2│ │Executor N│  │            │
│ │Partition1│ │Partition2│ │PartitionN│  │            │
│ │100M rows │ │100M rows │ │100M rows │  │            │
│ │          │ │          │ │          │  │            │
│ │Compute   │ │Compute   │ │Compute   │  │            │
│ │gradients │ │gradients │ │gradients │  │            │
│ └────┬─────┘ └────┬─────┘ └────┬─────┘  │            │
│      │            │            │        │            │
│      └────────────┼────────────┘        │            │
│                   │                     │            │
│                   ▼                     │            │
│           Aggregate Results             │            │
│           Update Model                  │            │
└──────────────────────────────────────────────────────┘

Each executor processes its partition in parallel
Gradients aggregated at driver for model update

Key Features

📊 DataFrame API

Works with Spark SQL, familiar SQL-like transformations

🔄 Pipeline API

Chain transformers & estimators like Scikit-learn

⚡ Distributed

Scales to terabytes, training parallelized across cluster

🎯 Production Features

Model persistence, versioning, batch scoring

Example: Logistic Regression on 500GB Data

from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler, StandardScaler
from pyspark.ml.classification import LogisticRegression
from pyspark.ml import Pipeline
from pyspark.ml.evaluation import BinaryClassificationEvaluator

# Initialize Spark with 100 executors
spark = SparkSession.builder \
    .appName("Large Scale ML") \
    .config("spark.executor.instances", "100") \
    .config("spark.executor.memory", "16g") \
    .config("spark.executor.cores", "4") \
    .getOrCreate()

# Load 500GB Parquet data (distributed across cluster)
df = spark.read.parquet("s3://data-lake/customer_features/")
# DataFrame: 5 billion rows × 50 features

# Train/test split (preserves distribution)
train_df, test_df = df.randomSplit([0.8, 0.2], seed=42)

# Feature engineering pipeline
assembler = VectorAssembler(
    inputCols=['feature_1', 'feature_2', ..., 'feature_50'],
    outputCol='features'
)

scaler = StandardScaler(
    inputCol='features',
    outputCol='scaled_features'
)

lr = LogisticRegression(
    featuresCol='scaled_features',
    labelCol='label',
    maxIter=10,
    regParam=0.01
)

# Build pipeline
pipeline = Pipeline(stages=[assembler, scaler, lr])

# Train (distributed across 100 executors)
model = pipeline.fit(train_df)

# Evaluate
predictions = model.transform(test_df)
evaluator = BinaryClassificationEvaluator(labelCol='label')
auc = evaluator.evaluate(predictions)
print(f"AUC: {auc:.3f}")

Result:
AUC: 0.892

Training Details:
• Data: 5 billion rows (500GB Parquet)
• Cluster: 100 executors × 16GB = 1.6TB total RAM
• Each executor: ~50M rows per partition
• Training time: ~25 minutes (10 iterations)
• Throughput: ~200M rows/minute

Example: Random Forest with Cross-Validation

from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.tuning import CrossValidator, ParamGridBuilder

# Random Forest model
rf = RandomForestClassifier(
    featuresCol='scaled_features',
    labelCol='label',
    seed=42
)

# Parameter grid for tuning
param_grid = ParamGridBuilder() \
    .addGrid(rf.numTrees, [50, 100, 200]) \
    .addGrid(rf.maxDepth, [5, 10, 15]) \
    .build()

# Cross-validator (3-fold)
cv = CrossValidator(
    estimator=rf,
    estimatorParamMaps=param_grid,
    evaluator=BinaryClassificationEvaluator(),
    numFolds=3,
    parallelism=10  # Train 10 models in parallel
)

# Train (tests 9 param combinations × 3 folds = 27 models)
cv_model = cv.fit(train_df)

# Best model
best_model = cv_model.bestModel
print(f"Best numTrees: {best_model.getNumTrees}")
print(f"Best maxDepth: {best_model.getOrDefault('maxDepth')}")

Result:
Best numTrees: 200 Best maxDepth: 10

Training Time: ~4 hours (27 models, each ~8 min)
Parallelism: 10 models trained simultaneously

When to Use Spark MLlib

Data exceeds 100GB (won't fit in single machine RAM)
Traditional ML algorithms (linear, trees, clustering)
Already using Spark for data processing
Need distributed training (not just inference)
Working with structured/tabular data at scale
Batch scoring on billions of records

Pros & Cons

✅ Pros

Scales to terabytes of data
Integrates with Spark ecosystem
Distributed training & inference
Familiar algorithms (trees, linear)
Pipeline API (like Scikit-learn)
Production features (versioning)

❌ Cons

Slower than Scikit-learn for small data
Limited algorithm selection
Requires cluster management
Debugging is harder (distributed)
No deep learning support
Steeper learning curve

TensorFlow: Production Deep Learning

Google's framework for scalable neural networks

TensorFlow is Google's deep learning framework, designed for production deployments. It excels at training massive neural networks on GPUs/TPUs and deploying them at scale. TensorFlow 2.x with Keras integration makes it much easier to use than v1.

Key Features

🚀 Production Ready

TensorFlow Serving, TFLite (mobile), TF.js (browser)

⚡ GPU/TPU Optimized

Automatic GPU acceleration, TPU support for massive scale

📊 Data Pipeline

tf.data API, efficient data loading from disk/cloud

🎯 Distributed Training

Multi-GPU, multi-node training with tf.distribute

Example: CNN for Image Classification (Keras API)

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Load dataset (60K training images)
(X_train, y_train), (X_test, y_test) = keras.datasets.cifar10.load_data()

# Normalize pixel values
X_train = X_train.astype('float32') / 255.0
X_test = X_test.astype('float32') / 255.0

# Build CNN model
model = keras.Sequential([
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.Flatten(),
    layers.Dense(64, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(10, activation='softmax')
])

# Compile
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# Train on GPU (automatic if available)
history = model.fit(
    X_train, y_train,
    epochs=20,
    batch_size=128,
    validation_split=0.2,
    verbose=1
)

# Evaluate
test_loss, test_acc = model.evaluate(X_test, y_test)
print(f"Test accuracy: {test_acc:.3f}")

Result:

Epoch 20/20
375/375 [======] - 3s 8ms/step - loss: 0.8234 - accuracy: 0.7145 - val_loss: 0.9012 - val_accuracy: 0.6923
Test accuracy: 0.694

Training Details:
• GPU: NVIDIA RTX 3090 (24GB)
• Time per epoch: ~3 seconds
• Total training: ~1 minute
• Speedup vs CPU: ~50x faster

Example: Efficient Data Pipeline (tf.data)

# Load images from directory (millions of files)
dataset = tf.data.Dataset.list_files('images/*/*.jpg')

def parse_image(filename):
    # Read and decode image
    image = tf.io.read_file(filename)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, [224, 224])
    image = image / 255.0
    
    # Extract label from path
    label = tf.strings.split(filename, '/')[-2]
    return image, label

# Build efficient pipeline
dataset = dataset \
    .map(parse_image, num_parallel_calls=tf.data.AUTOTUNE) \
    .cache() \
    .shuffle(buffer_size=10000) \
    .batch(32) \
    .prefetch(tf.data.AUTOTUNE)  # Prefetch next batch

# Train with pipeline (GPU never waits for data)
model.fit(dataset, epochs=10)

Optimization:
• Parallel image loading (8 CPU cores)
• Prefetching: CPU loads next batch while GPU trains
• Caching: First epoch reads disk, rest use RAM
• Result: GPU utilization ~95% (vs 60% without pipeline)

Example: Multi-GPU Training

# Distributed training across 4 GPUs
strategy = tf.distribute.MirroredStrategy()
print(f"Number of devices: {strategy.num_replicas_in_sync}")

# Build model inside strategy scope
with strategy.scope():
    model = keras.Sequential([
        layers.Dense(512, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(512, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(10, activation='softmax')
    ])
    
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

# Train (automatically distributed across GPUs)
model.fit(
    train_dataset,
    epochs=10,
    validation_data=val_dataset
)

Result:
Number of devices: 4

Performance:
• Effective batch size: 32 × 4 = 128
• Speedup: ~3.5x (4 GPUs with communication overhead)
• Training time: 10 epochs in 5 minutes (vs 18 min on 1 GPU)

When to Use TensorFlow

Production deep learning deployments
Need mobile/edge deployment (TFLite)
Training on TPUs (Google Cloud)
Large-scale distributed training
Image, speech, NLP tasks
When you need ecosystem (TensorBoard, Serving)

Pros & Cons

✅ Pros

Production-ready (Serving, Lite)
Excellent GPU/TPU support
Keras integration (easy API)
Great data pipeline (tf.data)
Multi-GPU/multi-node scaling
Strong ecosystem & community

❌ Cons

Steeper learning curve than PyTorch
Less Pythonic (more verbose)
Debugging can be harder
TF 1.x legacy code still exists
Slower research iteration
Overkill for simple models

PyTorch: Research & Dynamic Neural Networks

Facebook's framework for flexible deep learning

PyTorch is the most popular framework in research due to its Pythonic, intuitive API and dynamic computation graphs. It's easier to debug than TensorFlow and excels at NLP (transformers), computer vision, and reinforcement learning. Production deployment has improved significantly with TorchServe.

Key Features

🐍 Pythonic

Feels like NumPy, natural Python debugging with pdb

🔄 Dynamic Graphs

Define-by-run, change network architecture on the fly

🧪 Research First

Used by 70%+ of AI researchers, latest models available

⚡ Fast & Flexible

Excellent GPU support, easy to optimize custom layers

Example: Building a Custom Neural Network

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Define custom model (very Pythonic)
class CustomNN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(CustomNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x

# Initialize model, loss, optimizer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CustomNN(input_dim=100, hidden_dim=256, output_dim=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop (simple and explicit)
for epoch in range(10):
    model.train()
    for batch_X, batch_y in train_loader:
        batch_X, batch_y = batch_X.to(device), batch_y.to(device)
        
        # Forward pass
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

Result:

Epoch 1, Loss: 2.1234
Epoch 2, Loss: 1.8765
...
Epoch 10, Loss: 0.3421

Key Advantage: Explicit control, easy to debug with print(), pdb

Example: Transfer Learning (Pretrained ResNet)

import torchvision.models as models
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder

# Load pretrained ResNet50
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)

# Freeze all layers except final FC
for param in model.parameters():
    param.requires_grad = False

# Replace final layer for 10-class classification
model.fc = nn.Linear(model.fc.in_features, 10)
model = model.to(device)

# Only train the final layer
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

# Data augmentation
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Load dataset
train_dataset = ImageFolder('data/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# Fine-tune (fast, only training final layer)
for epoch in range(5):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Result:
Training time: ~10 minutes (vs hours training from scratch)
Accuracy: 94% (vs 87% random initialization)
Why fast: Only training 10 × 2048 = 20K params (vs 25M total)

Example: Distributed Data Parallel (Multi-GPU)

import os
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# Initialize distributed training
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)

# Wrap model in DDP
model = CustomNN(input_dim=100, hidden_dim=256, output_dim=10)
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

# Distributed sampler (each GPU gets different data)
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset
)
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    sampler=train_sampler
)

# Training loop (same as before!)
for epoch in range(10):
    train_sampler.set_epoch(epoch)
    for batch_X, batch_y in train_loader:
        # ... standard training code ...
        pass

Launch command:
torchrun --nproc_per_node=4 train.py

Result:
• 4 GPUs train in parallel
• Gradients synced across GPUs automatically
• Speedup: ~3.7x (with communication overhead)

When to Use PyTorch

Research and experimentation (fast iteration)
NLP tasks (transformers, Hugging Face)
Computer vision (torchvision models)
Custom architectures (dynamic graphs)
When you need flexibility and debugging
Academic projects and prototyping

Pros & Cons

✅ Pros

Most Pythonic & intuitive
Easy debugging (native Python)
Dynamic computation graphs
Huge research community
Excellent for NLP (transformers)
Fast iteration & experimentation

❌ Cons

Production deployment harder than TF
No built-in mobile support (yet)
Smaller deployment ecosystem
Less optimized for TPUs
More verbose than Keras
Steeper curve for beginners

Keras: High-Level Neural Network API

Simple, fast prototyping with minimal code

Keras is a high-level API that runs on top of TensorFlow (or PyTorch via keras-core). It's designed for fast experimentation with minimal code. Perfect for beginners and rapid prototyping, though less flexible than raw TensorFlow/PyTorch for custom architectures.

Example: Build & Train in 10 Lines

from tensorflow import keras
from tensorflow.keras import layers

# Build model (incredibly concise)
model = keras.Sequential([
    layers.Dense(128, activation='relu', input_shape=(784,)),
    layers.Dropout(0.2),
    layers.Dense(10, activation='softmax')
])

# Compile and train (one-liners)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(X_train, y_train, epochs=5, batch_size=32, validation_split=0.2)

# Evaluate
test_loss, test_acc = model.evaluate(X_test, y_test)
print(f"Test accuracy: {test_acc:.3f}")

Result: Working neural network in 10 lines!
Comparison: Same model in raw TensorFlow: ~40 lines, PyTorch: ~60 lines

When to Use Keras

Rapid prototyping (minimal code)
Learning deep learning (simplest API)
Standard architectures (CNNs, RNNs)
When development speed >flexibility
Teaching and educational projects

Pros & Cons

✅ Pros

Easiest to learn (beginner-friendly)
Most concise code
Fast prototyping
Excellent documentation
Backend-agnostic (TF, PyTorch)

❌ Cons

Less flexible than TF/PyTorch
Harder to customize layers
Abstracts away details (learning)
Limited for research

Decision Framework: Which Library?

Follow This Decision Tree:

1️⃣ What type of ML problem?

Traditional ML (trees, linear): → Scikit-learn or Spark MLlib
Deep Learning (neural networks): → Continue to question 2

2️⃣ Does your data fit in RAM?

YES (<100GB): → Scikit-learn (traditional) or TF/PyTorch/Keras (deep learning)
NO (>100GB tabular): → Spark MLlib
NO (>100GB images/text): → TF/PyTorch with data generators

3️⃣ For Deep Learning: Research or Production?

Research / Experimentation: → PyTorch (most flexible)
Production Deployment: → TensorFlow (better tooling)
Fast Prototyping: → Keras (easiest)

4️⃣ Specific Domain?

NLP (transformers): → PyTorch + Hugging Face
Computer Vision: → TensorFlow or PyTorch (both excellent)
Time Series: → Scikit-learn (small) or Spark MLlib (big)
Recommender Systems: → Spark MLlib (collaborative filtering)

Visual Summary

┌─────────────────────────────────────────────────────────────┐
│              ML FRAMEWORK DECISION TREE                     │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Data Size & Type?                                          │
│       │                                                     │
│       ├─ Small (<100GB), Tabular                            │
│       │    └─ Traditional ML → Scikit-learn ⭐              │
│       │                                                     │
│       ├─ Large (>100GB), Tabular                            │
│       │    └─ Distributed ML → Spark MLlib                  │
│       │                                                     │
│       └─ Any Size, Images/Text/Audio                        │
│            └─ Deep Learning?                                │
│                 │                                           │
│                 ├─ Research / Custom → PyTorch ⭐           │
│                 │                                           │
│                 ├─ Production / Scale → TensorFlow          │
│                 │                                           │
│                 └─ Fast Prototype → Keras ⭐                │
│                                                             │
└─────────────────────────────────────────────────────────────┘

Best Practices for ML at Scale

Always prototype on a small sample before using the full dataset or distributed framework.

# 1. Prototype on 1% sample with Scikit-learn
sample_df = full_df.sample(frac=0.01)
model = RandomForestClassifier()
model.fit(sample_X, sample_y)
# Result: Works? Accuracy good? → Continue

# 2. Scale to 10% on single machine
sample_df = full_df.sample(frac=0.10)
# Still fits in RAM? Use Scikit-learn

# 3. Only if needed, move to Spark MLlib
spark_df = spark.read.parquet("full_data/")
spark_model = SparkRandomForest()
spark_model.fit(spark_df)

⚠️ Anti-pattern: Starting with Spark for 10GB data. Scikit-learn would be 10x faster!

GPUs train 100x faster than CPUs, but they're often waiting for data. Use efficient data pipelines.

✅ Good Pipeline

Prefetch next batch while training
Parallel data loading (8+ workers)
Cache preprocessed data
Use data augmentation on GPU
Store data in binary format (HDF5, TFRecord)

❌ Bad Pipeline

Loading images from disk every epoch
Single-threaded data loading
CPU data augmentation (slow)
Keeping data in JSON/CSV
No prefetching (GPU waits)

For images/text, pretrained models save weeks of training time and often achieve better results.

Domain	Pretrained Model	Training Time Saved
Image Classification	ResNet-50, EfficientNet	Days → Hours
Object Detection	YOLO, Faster R-CNN	Weeks → Days
NLP	BERT, GPT, T5	Months → Hours
Speech	Wav2Vec, Whisper	Weeks → Days

Track the right metrics to catch problems early (overfitting, data leakage, convergence issues).

Essential Metrics to Track:

Training vs Validation Loss: Divergence = overfitting
Learning Rate: Too high = unstable, too low = slow
Gradient Norms: Exploding/vanishing gradients
GPU Utilization: Should be >90% during training
Throughput: Samples/second (catch bottlenecks)
Class Distribution: Imbalanced? Use weighted loss

💡 Tools: TensorBoard (TF), WandB (both), MLflow (tracking)

ML experiments are hard to reproduce. Version control for data and models is as important as code.

Code

Git + commit hash in logs

Data

DVC, date stamps, checksums

Models

MLflow, model registry

Common Pitfalls at Scale

❌ Data Leakage

Problem: Test data influences training (inflated metrics)

Solution: Split data BEFORE any preprocessing. Use time-based splits for time series. Never use test set for feature engineering decisions.

❌ Ignoring Class Imbalance

Problem: 99% negative class → 99% accuracy but useless model

Solution: Use class weights, oversample minority class, or use metrics like F1, AUC instead of accuracy.

❌ Training on Full Data Too Early

Problem: Waiting 10 hours only to find a bug

Solution: Always test on 1% sample first. Debug on small data, scale to full data only when everything works.

❌ Poor Hyperparameter Tuning

Problem: Manual tuning wastes time, grid search wastes compute

Solution: Use random search (faster than grid) or Bayesian optimization (Optuna, Hyperopt). Start with defaults from papers.

❌ Not Saving Checkpoints

Problem: Training crashes at epoch 9/10, lose everything

Solution: Save checkpoints every epoch. Keep best model by validation metric. Enable auto-resume for long training jobs.

❌ Premature Optimization

Problem: Optimizing code before proving model works

Solution: First make it work (simple model), then make it right (good performance), then make it fast (optimize if needed).

Real-World Example: E-Commerce Recommendation System

Problem: Recommend products to 50M users

Data: 5 years of purchase history (500GB), 10M products, 2B interactions

Approach:

Phase 1: Prototyping

Sample 1% data (5GB) → Scikit-learn
Test collaborative filtering baseline
Achieve 0.65 AUC in 2 hours
Validates approach works

Phase 2: Scaling

Full 500GB data → Spark MLlib ALS
Distributed training (100 executors)
Improve to 0.72 AUC
Training time: 4 hours

Phase 3: Deep Learning

Add neural collaborative filtering
PyTorch + embedding layers
Train on 4 GPUs
Final AUC: 0.78

Phase 4: Production

Export to TensorFlow Serving
Batch inference with Spark
Precompute recs nightly
Serve from Redis (sub-ms latency)

Complete Framework Comparison

Feature	Scikit-learn	Spark MLlib	TensorFlow	PyTorch	Keras
Best Use Case	Traditional ML	Big Data ML	Production DL	Research DL	Fast prototyping
Data Size Limit	~100GB	Unlimited (TB+)	Unlimited*	Unlimited*	Same as backend
Algorithms	100+ (trees, SVM)	20+ (basic ML)	Neural networks	Neural networks	Neural networks
GPU Support	❌ No	❌ No	✅ Excellent	✅ Excellent	✅ Via backend
Distributed	❌ No	✅ Yes (cluster)	✅ Multi-GPU/node	✅ Multi-GPU/node	✅ Via backend
Ease of Use	⭐⭐⭐⭐⭐	⭐⭐⭐	⭐⭐⭐	⭐⭐⭐⭐	⭐⭐⭐⭐⭐
Production Tools	Basic	Good	Excellent	Good	Via backend
Community	Huge	Large	Huge	Huge	Large
Learning Curve	Low	Medium	Medium-High	Medium	Low

* With proper data pipeline (tf.data, DataLoader)

Quick Reference: Choosing Your Stack

Structured/Tabular Data

< 100GB:

Scikit-learn ⭐
XGBoost, LightGBM

> 100GB:

Spark MLlib
H2O.ai

Images/Vision

Research:

PyTorch ⭐
+ torchvision

Production:

TensorFlow
+ TF Serving

Text/NLP

Transformers:

PyTorch ⭐
+ Hugging Face

Simple NLP:

Scikit-learn
+ spaCy

Key Takeaways

Start simple, Scikit-learn first, scale only when necessary
Data size drives choice, <100GB = single machine, >100GB = distributed
Spark MLlib for big tabular data, when data won't fit in RAM
PyTorch for research, most flexible, Pythonic, great for experimentation

TensorFlow for production DL, best deployment tools and ecosystem
Keras for prototyping, fastest way to test neural network ideas
Use transfer learning, pretrained models save weeks of training
Optimize data pipelines, don't let GPU wait for data

⚠️ Remember: The best framework is the one you know well. Master one before learning others. Most real projects use multiple frameworks at different stages (prototype → scale → deploy).

Further Learning Resources

📚 Books & Courses

Hands-On Machine Learning (Scikit-learn, TF)
Deep Learning with PyTorch
Spark: The Definitive Guide (MLlib chapter)
Fast.ai course (PyTorch-based)

🔗 Official Documentation

scikit-learn.org, Best ML docs
pytorch.org/tutorials, Excellent tutorials
tensorflow.org/guide, Comprehensive guides
spark.apache.org/mllib, MLlib reference

Bonus: End-to-End Recommendation Engine with PySpark ALS

Want to see everything from this lesson assembled into a working system? This open-source project puts PySpark ALS collaborative filtering, a FastAPI inference service, a React frontend, and a live event simulator together in one repo - covering the full ML lifecycle from raw interaction data to personalized recommendations served over HTTP.

The project covers:

PySpark ALS training - Trains separate article and product models from synthetic interaction data, outputting user/item factor parquets for fast dot-product inference
FastAPI inference API - Loads parquet factors at startup, computes numpy dot-product recommendations per request, and hot-reloads whenever the model files change
Live simulator - Continuously generates interaction events, seeds them into PostgreSQL, and triggers model retraining on a schedule - demonstrating a real ML feedback loop
AWS deployment guide - Maps every component to managed services: Glue for PySpark training, Lambda + Mangum for inference, RDS + RDS Proxy for the database, S3 for model storage, and ECS Fargate for the simulator
Production-ready alternatives - Discusses when to swap ALS for theimplicit library, dot-product for FAISS, on-demand inference for batch pre-computation into DynamoDB, and how to add a two-stage retrieval + ranking pipeline

Explore the project:gitlab.com/alek.cora.glez/recommendation-engine-platform

Back to "Data Analytics at Scale"

Big Data FundamentalsLesson 10

Next: "Cloud Big Data Services"

Machine Learning on Big Data

When Your Data Won't Fit in RAM

The Scale Problem: Data vs Model Complexity

📊 Big Data, Simple Model

🧠 Small Data, Complex Model

ML Framework Landscape

Scikit-learn: The Gold Standard for Traditional ML

Key Strengths

🎯 Consistent API

📚 Rich Algorithms

🔧 Built-in Tools

⚡ Production Ready

Example: Classification with Random Forest

Example: Hyperparameter Tuning

When to Use Scikit-learn

Pros & Cons

✅ Pros

❌ Cons

Spark MLlib: Distributed Machine Learning

Architecture: Distributed Training

Key Features

📊 DataFrame API

🔄 Pipeline API

⚡ Distributed

🎯 Production Features

Example: Logistic Regression on 500GB Data

Example: Random Forest with Cross-Validation

When to Use Spark MLlib

Pros & Cons

✅ Pros

❌ Cons

TensorFlow: Production Deep Learning

Key Features

🚀 Production Ready

⚡ GPU/TPU Optimized

📊 Data Pipeline

🎯 Distributed Training

Example: CNN for Image Classification (Keras API)

Example: Efficient Data Pipeline (tf.data)

Example: Multi-GPU Training

When to Use TensorFlow

Pros & Cons

✅ Pros

❌ Cons

PyTorch: Research & Dynamic Neural Networks

Key Features

🐍 Pythonic

🔄 Dynamic Graphs

🧪 Research First

⚡ Fast & Flexible

Example: Building a Custom Neural Network

Example: Transfer Learning (Pretrained ResNet)

Example: Distributed Data Parallel (Multi-GPU)

When to Use PyTorch

Pros & Cons

✅ Pros

❌ Cons

Keras: High-Level Neural Network API

Example: Build & Train in 10 Lines

When to Use Keras

Pros & Cons

✅ Pros

❌ Cons

Decision Framework: Which Library?

Follow This Decision Tree:

1️⃣ What type of ML problem?

2️⃣ Does your data fit in RAM?

3️⃣ For Deep Learning: Research or Production?

4️⃣ Specific Domain?

Visual Summary

Best Practices for ML at Scale

1. Start Small, Scale Later

2. Optimize Data Loading (Don't Starve Your GPU)

✅ Good Pipeline

❌ Bad Pipeline

3. Use Transfer Learning (Don't Train From Scratch)

4. Monitor Training with Good Metrics

Essential Metrics to Track:

5. Version Everything (Data, Code, Models)

Code