import numpy as np
from sklearn.datasets import make_classification
from torch import nn
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from skorch import NeuralNetClassifier
from skorch.callbacks import EpochScoring, TensorBoard

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

class ClassifierModule(nn.Module):
    def __init__(
            self,
            num_units=10,
            nonlin=F.relu,
            dropout=0.5,
    ):
        super().__init__()
        self.dense0 = nn.Linear(20, num_units)
        self.nonlin = nonlin
        self.dropout = nn.Dropout(dropout)
        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 = F.relu(self.dense1(X))
        X = F.softmax(self.output(X), dim=-1)
        return X

writer = SummaryWriter()
net = NeuralNetClassifier(
    ClassifierModule,
    max_epochs=20,
    lr=0.1,
    callbacks=[
        EpochScoring(scoring='roc_auc', lower_is_better=False),
        TensorBoard(writer),
    ],
)
net.fit(X, y)