Can I resume training from checkpoints using the ktrain library?

Viewed 1812

ktrain is a lightweight wrapper for the deep learning library TensorFlow Keras (and other libraries) to help build, train, and deploy neural networks and other machine learning models. Can I resume training from checkpoints using the ktrain library?

1 Answers

Yes, you can. This is answered in the ktrain FAQ. I will copy the answer here:

Method 1: Using Predictor API (works for any model)

# save model and Preprocessor instance after partially training
ktrain.get_predictor(model, preproc).save('/tmp/my_predictor')

# reload Predictor and extract model
model = ktrain.load_predictor('/tmp/my_predictor').model

# re-instantiate Learner and continue training
learner = ktrain.get_learner(model, train_data=trn, val_data=val)
learner.fit_onecycle(2e-5, 1)

Note that preproc here is a Preprocessor instance. If using a data-loading function like texts_from_csv or images_from_folder, it will be the third return value from the function. Or, if using the Transformer API for text classification, it will be the output of invoking text.Transformer (i.e., preproc = text.Transformer('bert-base-uncased', ...)).

Method 2: Using transformers library (if training Hugging Face Transformers model)

If the model is a Hugging Face transformers model, you can use transformers directly:

# save model using transformers API after partially training
learner.model.save_pretrained('/tmp/my_model')

# reload the model using transformers directly
from transformers import *
model = TFAutoModelForSequenceClassification.from_pretrained('/tmp/my_model')
model.compile(loss='categorical_crossentropy',optimizer='adam', metrics=['accuracy'])

# re-instantiate Learner and continue training
learner = ktrain.get_learner(model, train_data=trn, val_data=val)
learner.fit_onecycle(2e-5, 1)

Method 3: Using checkpoint_folder argument to save model weights

The checkpoint_folder argument (e.g., learner.autofit(1e-4, 4, checkpoint_folder='/tmp/saved_weights')), saves the weights only of the model after each epoch. The weights of any epoch can be reloaded into the model using the model.load_weights method as you normally would in tf.Keras. You just need to first re-create the model first. For instance, if training an NER model, it would work as follows:

# recreate model from scratch
import ktrain
from ktrain import text
model = text.sequence_tagger(...
# load checkpoint weights from 3rd epoch into model
model.load_weights('../models/checkpoints/weights-03.hdf5')
# recreate learner
learner = ktrain.get_learner(model, ...
# continue training here

Finally, there is also a learner.save_model and learner.load_model methods intended for saving and reloading models when training interactively during a single session.

Related