CODEWITHSUNDEEP
pythontensorflownlphuggingface-transformers
shao_chen
shao_chen

Reputation: 29

How to save the best model of each epoch with transformers bert in tensorflow

I uesd TFBertModel and Tensorflow model to combined and training with Hugging Face transformers.I want to save the best model of val_accuracy of each epoch.I used 'tensorflow checkpoint', but I got the error.How do I save the best model of each epoch with transformers bert in tensorflow?

from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.initializers import TruncatedNormal
from tensorflow.keras.losses import CategoricalCrossentropy,BinaryCrossentropy
from tensorflow.keras.metrics import CategoricalAccuracy,BinaryAccuracy
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.utils import plot_model
from transformers import AutoTokenizer,TFBertModel
import tensorflow as tf
from tensorflow.keras.layers import Input, Dense
tokenizer = AutoTokenizer.from_pretrained('bert-large-uncased')
bert = TFBertModel.from_pretrained('bert-large-uncased')
max_len = max_length
input_ids = Input(shape=(max_len,), dtype=tf.int32, name="input_ids")
input_mask = Input(shape=(max_len,), dtype=tf.int32, name="attention_mask")
# embeddings = dbert_model(input_ids,attention_mask = input_mask)[0]


embeddings = bert(input_ids,attention_mask = input_mask)[1] #(0 is the last hidden states,1 means pooler_output)
# out = tf.keras.layers.GlobalMaxPool1D()(embeddings)
out = tf.keras.layers.Dropout(0.1)(embeddings)

out = Dense(128, activation='relu')(out)
out = tf.keras.layers.Dropout(0.1)(out)
out = Dense(32,activation = 'relu')(out)

y = Dense(1,activation = 'sigmoid')(out)
    
model = tf.keras.Model(inputs=[input_ids, input_mask], outputs=y)
model.layers[2].trainable = True
#model.save_weights('path/savefile')

Upvotes: 0

Views: 2952

Answers (1)

Gerry P
Gerry P

Reputation: 8092

if you are using tensorflow then you can create the callback below which will save the model for each epoch. filepath is the path to the directory where you want to save your model. model is the name of your compiled model. The models will be saved in the format epoch-validation loss.h5

class model_per_epoch(keras.callbacks.Callback):
    def __init__(self, model,filepath):
        self.filepath=filepath
        self.model=model
    def on_epoch_end(self,epoch, logs=None):
        v_loss=logs.get('val_loss') 
        name= str(epoch) +'-' + str(v_loss)[:str(v_loss).rfind('.')+3] + '.h5'
        file_id=os.path.join(self.filepath, name)
        self.model.save(file_id)
save_dir=r'c:\Temp\models'
callbacks=[model_per_epoch(model, save_dir)]

in model.fit include callback=callbacks . Make sure the directory you are saving the model to exists. The code below is a more sophisticated version of the callback. An additional parameter save_best_only is added. If set to True, only the model with the lowest validation loss is saved. Also the model is loaded with the weights for the epoch with the lowest validation loss so you can use your model directly to do predictions without having to load the saved model. This can be a significant time saver.

class model_per_epoch(keras.callbacks.Callback):
    def __init__(self, model,filepath,save_best_only):
        self.filepath=filepath
        self.model=model
        self.save_best_only=save_best_only
        self.lowest_loss=np.inf
        self.best_weights=self.model.get_weights()
    def on_epoch_end(self,epoch, logs=None):
        v_loss=logs.get('val_loss')
        if v_loss< self.lowest_loss:
            self.lowest_loss =v_loss
            self.best_weights=self.model.get_weights()
            self.best_epoch=epoch +1
        if self.save_best_only==False:
            name= str(epoch) +'-' + str(v_loss)[:str(v_loss).rfind('.')+3] + '.h5'
            file_id=os.path.join(self.filepath, name)
            self.model.save(file_id)
    def on_train_end(self, logs=None):
        if self.save_best_only == True:
            self.model.set_weights(self.best_weights)
            name= str(self.best_epoch) +'-' + str(self.lowest_loss)[:str(self.lowest_loss).rfind('.')+3] + '.h5'
            file_id=os.path.join(self.filepath, name)
            self.model.save(file_id)
            print(' model is returned with best weiights from epoch ', self.best_epoch)
            
save_dir=r'c:\Temp\models'
save_best_only= True
callbacks=[model_per_epoch(model, save_dir, save_best_only)]

Upvotes: 1

Related Questions