Training a model

Below, we define our own PyTorch Module and train it on a toy classification dataset using skorch NeuralNetClassifier:

import numpy as np
from sklearn.datasets import make_classification
from torch import nn

from skorch import NeuralNetClassifier

X, y = make_classification(1000, 20, n_informative=10, random_state=0)
X = X.astype(np.float32)
y = y.astype(np.int64)

class MyModule(nn.Module):
    def __init__(self, num_units=10, nonlin=nn.ReLU()):
        super(MyModule, self).__init__()

        self.dense0 = nn.Linear(20, num_units)
        self.nonlin = nonlin
        self.dropout = nn.Dropout(0.5)
        self.dense1 = nn.Linear(num_units, 10)
        self.output = nn.Linear(10, 2)

    def forward(self, X, **kwargs):
        X = self.nonlin(self.dense0(X))
        X = self.dropout(X)
        X = self.nonlin(self.dense1(X))
        X = self.output(X)
        return X

net = NeuralNetClassifier(
    # Shuffle training data on each epoch

net.fit(X, y)
y_proba = net.predict_proba(X)


In this example, instead of using the standard softmax non-linearity with NLLLoss as criterion, no output non-linearity is used and CrossEntropyLoss as criterion. The reason is that the use of softmax can lead to numerical instability in some cases.

In an sklearn Pipeline

Since NeuralNetClassifier provides an sklearn-compatible interface, it is possible to put it into an sklearn Pipeline:

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

pipe = Pipeline([
    ('scale', StandardScaler()),
    ('net', net),

pipe.fit(X, y)
y_proba = pipe.predict_proba(X)

What’s next?

Please visit the Tutorials page to explore additional examples on using skorch!