Model not learning after upgrade to tensorflow 2 and version 1.0.0

[digital-thinking]

I have just upgraded from TensorFlow 1 with efficientnet 0.0.4 to the TensorFlow 2 version of the model. The model looks identical (besides the names of the tensors), params are equal, training data is equal. However, with TensorFlow 2.0 the model does not learn anymore. I have tried TensorFlow 2.0 with CUDA 10.0 and TensorFlow 2.1 with CUDA 10.1 both with the same results. I wonder if anyone else has also noticed issues like that. With TF1 and the 0.0.4 model everything works fine. Also the training is slower than before. (RTX 2080Ti)

TF2(efficientnet 1.0.0)

from efficientnet.tfkeras import EfficientNetB5
base_model = EfficientNetB5(weights='imagenet', include_top=False)
x = base_model.output
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dense(len(class_weights), activation='softmax')(x)
model = tf.keras.models.Model(inputs=base_model.get_input_at(0), outputs=x)
model.compile(optimizer=tf.keras.optimizers.Adam(lr=0.01), loss='categorical_crossentropy',  metrics=['accuracy'])

TF1(efficientnet 0.0.4):

from efficientnet import EfficientNetB5
base_model = EfficientNetB5(include_top=False)
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(len(class_weights), activation='softmax')(x)
model = Model(inputs=base_model.get_input_at(0), outputs=x)
model.compile(optimizer=optimizers.Adam(lr=0.01), loss='categorical_crossentropy',  metrics=['accuracy'])

[TomLin]

I've also run into the same problem. Had no other way out but to revert back to 0.0.4. Perhaps there's a bug in the new version of pre-train weights.

[das-sein]

I thought I had this problem, then discovered it was simply my model setup. I'm training successfully on version 1.1.0.

[MarcRoigCampos]

Can you show which setup are you using now? 🙏

[das-sein]

@MarcRoigCampos Please forgive the untidiness of this, it's being used for a class assignment and wasn't really meant to be seen. https://gist.github.com/latrare/773128d9a5f15e31854ea196b414b89a

Note, I am not using the pre-defined weights here, so that maybe where this issue lies if even after adjusting your program it still does not train the model.

[MarcRoigCampos]

@latrare Thanks a lot!!