tensorflow_transfer_learning_example.py

import tensorflow as tf
import tensorflow_hub as hub

# colelct the feature extractor of mobile net
feature_extractor_url = "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/2" #@param {type:"string"}
feature_extractor_layer = hub.KerasLayer(feature_extractor_url,
                                         input_shape=(IMG_HEIGHT, IMG_WIDTH,3))

feature_extractor_layer.trainable = False

# Build the model
model = tf.keras.Sequential([
  feature_extractor_layer,
  tf.keras.layers.Dense(len(CLASS_NAMES))
])

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

log_dir = "logs\\fit\\" + 'tlmnetv2_' + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir)

history = model.fit(
train_data_gen,
steps_per_epoch=image_count_train // BATCH_SIZE,
epochs=EPOCHS,
validation_data=val_data_gen,
validation_steps=image_count_validation // BATCH_SIZE,
callbacks=[tensorboard_callback]
)