Save the best performing model (but not regarding validation acc) in Pytorch

user_12201107 0
Viewed 406

Is there a way to replicate behavior of Keras earlyStopping method in PyTorch? For example, if I want best model in regards of f1 score rather than acquired validation accuracy, is there an equivalent to Keras earlyStopping = EarlyStopping(monitor='f1_score', patience=10, verbose=0, mode='min')?

pytorch keras machine-learning
1 Answers

Pytorch has no callback functions like Keras does, so you'd have to take care of this manually.

Here's a rough pseudo-code outline:

NUM_EPOCHS = 100
train_data, val_data = ...
model = ...
loss_fn = ...
optimizer = ...

# for early stopping
curr_best_f1 = 0.0
patience = 10

for epoch in range(NUM_EPOCHS):
    
    # training step
    model.train()
    train_step(model, loss_fn, optimizer, train_data)

    # validation step
    model.eval()
    val_preds, val_labels = val_step(model, loss_fn, val_data)
    
    # check early stopping criterion
    f1_score = F1_Score(val_preds, val_labels)
    if f1_score >= curr_best_f1:
        curr_best_f1 = f1_score
        patience = 10
    else:
        patience -= 1
    
    if patience == 0:
        break

Change F1_Score to whatever metric you like to use and the >= accordingly if you like to have min instead of max. You may also want to put the early stopping checking in a separate function or small class.

user_112208840
Related