How to tune number of epochs?

[jajupmochi]

Hi, may I ask how I can find best number of epochs (not max_epochs) through cross validation? Is it possible to do something like:

from sklearn.model_selection import GridSearchCV, cross_validate
from skorch import NeuralNetClassifier

net = NeuralNetClassifier(...)
params = {
	'n_epochs': [10, 20, 30, 40, 50],
}
gs = GridSearchCV(
	estimator=net,
	param_grid=params,
	refit=True,
	cv=inner_cv,
	scoring='accuracy',
)
test_predictions = cross_validate(
	gs,
	X,
	y,
	cv=outer_cv,
	scoring=scoring,
	verbose=0,
	n_jobs=n_cores_cv
)
...

If it is possible, may I ask how the models are evaluated on each n_epochs? Is an individual model fitted for each n_epochs, or the valid metrics are calculated and recorded along the training epochs of a single model, so that the model will be fitted only once for one split of dataset?

Thanks!

[BenjaminBossan]

Hi, I converted the issue into a discussion, I hope you don't mind.

may I ask how I can find best number of epochs (not max_epochs)

Could you explain what the difference is supposed to be? max_epochs is the number of epochs that the model is trained, so it sounds like the same as what you suggest.

Is an individual model fitted for each n_epochs, or the valid metrics are calculated and recorded along the training epochs of a single model, so that the model will be fitted only once for one split of dataset?

Using GridSearchCV and cross_validate as you show, it would train one copy of the model per hyper-parameter and per split. This is purely sklearn functionality, not skorch.

If you don't want a separate model for each epoch count, you could just use the highest one. The net stores its history in the net.history attribute, which you can use, once training is finished, to check it's performance at different stages of the training. For you to get access to each net (one for each split), you would have to return it from cross_validate by passing return_estimator=True. Note that this will take up extra memory.

Be aware that the net does not "know" what scoring argument you pass to grid search. If you have additional metrics in there that you want the net to store in its history, you should add a scoring callback with that metric to the net.

[jajupmochi]

Thx @BenjaminBossan! I am trying to tune the n_epochs with a cross-validation, like this:

import numpy as np
from sklearn.metrics import accuracy_score
from sklearn.model_selection import StratifiedKFold, ParameterGrid
import torch
import torch.nn as nn
from skorch import NeuralNetClassifier
from skorch.callbacks import EpochScoring

# Define your custom scoring function
def custom_accuracy(net, dataset, y_true):
    y_pred = net.predict(dataset)
    return accuracy_score(y_true, y_pred)

# Define a simple neural network using PyTorch
class SimpleNN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, X):
        X = self.fc1(X)
        X = self.relu(X)
        X = self.fc2(X)
        return X

max_epochs = 1000

param_grid = {
    'n_epochs': list(range(10, max_epochs + 1)),  # Search from 10 to max_epochs
}

param_list = list(ParameterGrid(param_grid))

kf = StratifiedKFold(n_splits=5, shuffle=True)

# Initialize a dictionary to store scores for each n_epochs
epoch_scores = {}

for params in param_list:
    # Initialize a list to collect accuracy scores for each split
    split_scores = []

    # for each inner CV fold:
    for train_index, valid_index in kf.split(dataset_app, y_app):
        net = NeuralNetClassifier(
            module=SimpleNN,
            ...
            max_epochs=params['n_epochs'],
            callbacks=[
                ('accuracy', EpochScoring(custom_accuracy, name='accuracy', lower_is_better=False, on_train=False)),
            ],
            ...
        )

        net.fit(dataset_app[train_index], y_app[train_index], valid_dataset=dataset_app[valid_index], valid_y=y_app[valid_index])

        # Append the accuracy score to the list
        split_scores.append(net.history[:, 'accuracy'])

    # Calculate the average accuracy score over the splits
    avg_score = np.mean(split_scores, axis=0)
    
    # Store the average score for this n_epochs
    epoch_scores[params['n_epochs']] = avg_score

# Find the best n_epochs based on the highest average score
best_n_epochs = max(epoch_scores, key=lambda k: epoch_scores[k][-1])
print(f"Best number of epochs: {best_n_epochs}")
...

The code records accuracy scores for each n_epochs in each split, averages the scores over the splits for each n_epochs, and finally finds the n_epochs with the highest average score. In this way, I only need to run one model for each split, as you suggested. You can see here the difference between what I meant n_epochs and max_epochs.

However, I have also other hyperparameters to tune together with n_epochs, such as feats_dim and n_layers, which require a separate model for each hyperparam setting.

My question is: is there a "simple" way to do this with skorch, ideally compatible with sklearn, as in my first question?

Thx!

[jajupmochi]

@BenjaminBossan I also put this question on stack overflow and got this answer:

from sklearn.model_selection import GridSearchCV, cross_validate
from skorch import NeuralNetClassifier
from skorch.callbacks import EpochScoring

# Define your neural network module
class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, X):
        # Define your forward pass logic here
        pass

net = NeuralNetClassifier(
    module=MyModule,
    optimizer=torch.optim.Adam,
    criterion=nn.CrossEntropyLoss,
    max_epochs=50,  # Set a maximum number of epochs for the grid search
    callbacks=[
        ('accuracy', EpochScoring('accuracy', lower_is_better=False)),
    ]
)

params = {
    'module__n_epochs': [10, 20, 30, 40, 50],
}

gs = GridSearchCV(
    estimator=net,
    param_grid=params,
    refit=True,
    cv=inner_cv,
    scoring='accuracy',
)

test_predictions = cross_validate(
    gs,
    X,
    y,
    cv=outer_cv,
    scoring=scoring,
    verbose=0,
    n_jobs=n_cores_cv
)

If I understand you correctly, this code will run a separate model for each epoch count, right?

[BenjaminBossan]

So from what I understand, n_epochs is not really a hyper-parameter in your case, because you're not really trying to tune it. Instead, what you're interested in is knowing at what epoch the score is best, given a set of hyper-parameters.

The SO answer is not correct from what I can tell. It tries to pass the n_epochs argument to the module, but 1) the module does not accept that argument and 2) the module wouldn't know what to do with it.

In theory, ideally, you could run sklearn's GridSearchCV with the hyper-parameters you're interested in (so the other stuff, not max_epochs), get back the net, inspect the history, and determine at what epoch the score is optimal. But that doesn't work because GridSearchCV does not allow you to return each net. This is a limitation of sklearn.

However, what you can do is the following: You define your parameter grid with the other hyper-params. You use the EarlyStopping callback from skorch in the net and set a very high value for max_epochs -- also, ensure that EarlyStopping monitors the score you're interested in. You fit an sklearn GridSearchCV with (with refit=True, which is the default). Sklearn will figure out the best hyper-parameters and skorch will stop the training early when the score stops improving. After the grid seach has finished fitting, sklearn refits the net based on the best hyperparameters. You can find that net as gs.best_estimator_. There, you can check the history to see when the best score was reached. This should give you the information you need and at the same time by quite fast, as you don't have to fit one model per max_epochs value.