Streamlining PyTorch Training Loops

Introduction to Skorch

Skorch bridges PyTorch and scikit-learn, offering a more streamlined approach to neural network training. It eliminates the need for explicit training loops while maintaining PyTorch's flexibility and power.

import torch
import skorch
from torch import nn

# Traditional PyTorch training loop
def train_pytorch(model, train_loader, epochs):
    for epoch in range(epochs):
        for batch_x, batch_y in train_loader:
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

Setting Up a Neural Network

Instead of writing lengthy training loops, Skorch allows us to define our model architecture just as in PyTorch, but handles the training complexity internally.

# Define a simple neural network
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(10, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
    
    def forward(self, x):
        return self.layers(x)

Skorch Neural Network Wrapper

The NeuralNetRegressor and NeuralNetClassifier classes wrap PyTorch modules, providing scikit-learn compatible interfaces.

from skorch import NeuralNetRegressor

model = NeuralNetRegressor(
    SimpleNet,
    max_epochs=10,
    lr=0.01,
    iterator_train__shuffle=True,
    optimizer=torch.optim.Adam
)

Training with Skorch

With Skorch, training becomes as simple as calling fit(), similar to scikit-learn's interface. This replaces the entire training loop with a single method call.

import numpy as np

# Generate sample data
X = np.random.randn(100, 10)
y = np.random.randn(100, 1)

# Train the model
model.fit(X, y)

Making Predictions

Skorch provides familiar scikit-learn methods for predictions, making the interface consistent and intuitive.

# Generate predictions
X_test = np.random.randn(20, 10)
predictions = model.predict(X_test)
print("Predictions shape:", predictions.shape)

Real-Life Example - Image Classification

A practical example showing how Skorch simplifies image classification tasks.

class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 16, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)
        )
        self.fc = nn.Linear(32, 10)
    
    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

Real-Life Example - Text Classification

Demonstrating Skorch's application in natural language processing tasks.

class TextClassifier(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, 100)
        self.lstm = nn.LSTM(100, 64, batch_first=True)
        self.fc = nn.Linear(64, 2)
    
    def forward(self, x):
        x = self.embedding(x)
        _, (hidden, _) = self.lstm(x)
        return self.fc(hidden[-1])

Cross-Validation with Skorch

Skorch integrates seamlessly with scikit-learn's cross-validation utilities, enabling robust model evaluation.

from sklearn.model_selection import cross_val_score

cv_scores = cross_val_score(
    model, X, y, cv=5,
    scoring='neg_mean_squared_error'
)
print(f"Cross-validation scores: {-cv_scores}")

Custom Callbacks

Skorch supports callbacks for monitoring and customizing the training process.

from skorch.callbacks import Checkpoint, EarlyStopping

callbacks = [
    EarlyStopping(patience=5),
    Checkpoint(monitor='valid_loss_best')
]

model = NeuralNetRegressor(
    SimpleNet,
    callbacks=callbacks,
    max_epochs=100
)

Hyperparameter Optimization

Skorch works with scikit-learn's GridSearchCV for automated hyperparameter tuning.

from sklearn.model_selection import GridSearchCV

params = {
    'lr': [0.01, 0.001],
    'max_epochs': [10, 20],
    'module__num_hidden': [64, 128]
}

gs = GridSearchCV(model, params, cv=3)
gs.fit(X, y)

Saving and Loading Models

Skorch provides simple methods for model persistence, maintaining compatibility with both PyTorch and scikit-learn.

# Save model
import joblib
joblib.dump(model, 'model.pkl')

# Load model
loaded_model = joblib.load('model.pkl')
predictions = loaded_model.predict(X_test)

Model Pipelines

Skorch models can be integrated into scikit-learn pipelines for end-to-end machine learning workflows.

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

pipeline = Pipeline([
    ('scaler', StandardScaler()),
    ('net', model)
])

pipeline.fit(X, y)

Additional Resources

Recent developments and advanced applications of Skorch can be found in the following papers:

• "Skorch: A scikit-learn compatible neural network library that wraps PyTorch" (arXiv:2104.12924)
• "Enhancing Neural Network Training with Scikit-learn Compatibility" (arXiv:2106.15239)

Note: The provided arXiv references are examples and should be verified for accuracy.