I trained a model with the following architecture:
bert_config = BertConfig.from_pretrained(MODEL_NAME) bert_config.output_hidden_states = True backbone = TFAutoModelForSequenceClassification.from_pretrained(MODEL_NAME,config=bert_config) input_ids = tf.keras.layers.Input(shape=(MAX_LENGTH,), name='input_ids', dtype='int32') features = backbone(input_ids)[-1] pooling = tf.keras.layers.GlobalAveragePooling1D()(features) dense = tf.keras.layers.Dense(len(label2id), name='output',activation=tf.nn.softmax)(pooling) model = tf.keras.Model(inputs=[input_ids], outputs = [dense])
and saved model in different ways. The first one is
and the second one is
So, I am able to load my model (
model.load_weights()) from both of these options without any error. Moreover, inference is fine. In short, everything works as I expect.
But if I start a new session and load my model again then inference is bad, like model has random weights instead of my saved weights.
I was trying other options of saving models as well, but they do not work also. Probably there is a special way to save a model with transformer layer?
Thanks in advance!