Deep Learning with Best-so-far instead of Where-you-are

Question

It is my understanding that when training a Deep NN in Tensorflow/PyTorch/... we only keep the current state of the network in memory, except perhaps when we manually decide to save the current weights to the HDD/SSD.

Now, naively speaking it may seem reasonable to not only remember the current state (i.e. the current values of the trained weights), i.e. "where we are", but also the "best" weights "so far" by some metric, such as the validation error. Immediately, this approach doubles our memory requirements, especially if we want to keep everything in the GPU memory.

Is this done in practice? If not what are arguments against it?

---

Now there are certainly details that make this question more tricky. If you saved only each epoch it's probably fine to just save your weights to the HDD without significant loss in speed. If you remember the best weights for each step of SGD than you would have to compute the validation error over and over, which is costly. Or you use perhaps an estimate of the training error based on its gradient, which however could be leading to overfitting if you are not careful. You could also mix and match, or you could compute the validation error only on a subset of the data, or only remember the best every $k$ steps, etc. etc. You might also want to go back and restart from a "better" set of weights in the hopes of getting a better trajectory.

I also don't know whether in practice at the of training people typically, deliberately, go back to a previous "checkpoint" they might have saved somewhere.

Answer 1

Is this done in practice?

Yes, this is done normally when using (lack of) improvements to validation metrics as a stop criterion, and many libraries support it as standard. Depending on the library, you may find you need to add a little code to keep a copy of the best-so-far weights, but some will do it automatically by default, or based on setting params on the train or fit function.

For example, Keras' EarlyStopping class has a restore_best_weights parameter. Using this class in your main fit function, and setting the param to true will do what you want automatically with no other code required.

If not what are arguments against it?

Over-fitting to the validation set is a possible concern, as running the validation checks 100s of times to decide the "best" model may lead to some maximisation bias, and make decisions between other hyperparameters than the number of epochs less reliable.

Answer 2

Common practice

Model checkpoints are often saved to the HDD to keep the GPU memory free.

At every epoch, a selection metric (e.g. validation loss) is evaluated after the training stage completes, and we save the:

• Top-k best models (by model selection metric)
• Latest model at end of epoch

Example training loop:

best_val_loss = np.inf

for epoch in range(num_epochs):
    out_train = train_epoch(model)
    out_val = validate_epoch(model)

    save(model, "last.ckpt")

    if out_val["loss"] < best_val_loss:
        save(model, "best.ckpt")
        best_val_loss = out_val["loss"]

It is possible that there exist models that best minimize the selection metric during an epoch. However, since the gain is likely quite minimal, it is much more practical to choose the strategy mentioned above.

---

Is there something better?

You propose some interesting strategies. It is possible that there might be a better process than the typical "train/validation/test over shuffled epoch with batches", but it needs to be shown that such a process leads to noticeable improvements. The convenience and comprehensibility of using a common approach should not be underestimated.

You might also want to go back and restart from a "better" set of weights in the hopes of getting a better trajectory.

Restarting from a better model is almost like training using the validation set, though admittedly not as severe as using validation gradients to optimize the model directly. Typical problems to consider when using such an approach are overfitting and difficulty in escaping local minima.