Feature Request in Layers: Store best model with save_best_only.

[cpury]

TensorFlow.js version

0.10.3

Browser version

Latest Chrome

Describe the problem or feature request

When training on a just a few samples as usually happens in the browser, training and validation loss tend to jump around a bit.

It would be great if we could automatically get the best model after training finishes.

This is possible in Keras with the ModelCheckpoint(save_best_only=True) callback and afterwards loading the weights from that saved file.

Alternatively, if I knew how to copy a model in its current state, it would be straightforward to write the logic for this in a onEpochEnd callback.

[bileschi]

Great suggestion! FYI, model copying is coming soon. Please follow tensorflow/tfjs-core#1038 for guidance on that front.

[cpury]

@bileschi it seems model copying is now implemented, which indirectly makes this feature possible.

Is there a best practice / guide to retain the best model in training? That would also be a useful addition to the docs, I believe.

[bileschi]

There is not such a guide or demo at this point. Agreed that it would make an excellent addition to the tutorials.

[cpury]

@bileschi I went ahead and tried it out. Works great! Thanks for pointing me to the other issue.

Here's what I did:

let bestValLoss = Number.MAX_SAFE_INTEGER;
const bestModelPath = 'localstorage://best-model';

model.fit(..., {
    ...,
    callbacks: {
        onEpochEnd: async function (epoch, logs) {
            if (logs.val_loss < bestValLoss) {
                // Store best model:
                bestValLoss = logs.val_loss;
                await model.save(bestModelPath);
            }
        },
        onTrainEnd: async function () {
            // Load best model:
            model = await tf.loadModel(bestModelPath);
        },
    }
});

Note that if you want to keep training afterwards, you need to .compile() the model again.

[cpury]

Quick question: Does TF.JS await the callbacks? Otherwise there might be a race condition here.

[google-ml-butler]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs. Thank you.

[google-ml-butler]

Closing as stale. Please @mention us if this needs more attention.