Things to be noted when using create_model(model_url, num_classes)

[PengJuan-AI]

So I try to use models created in 04-transfer-learning-in-tensorflow-part1-feature-extraction using function create_model(model_url, num_classes) but ran into this problem:

ValueError: logitsandlabels must have the same shape, received ((None, 2) vs (None, 1)).

[PengJuan-AI]

It turns out that in function create_model, you have to make some changes to make it adapt to binary classification, which is shown below:

def create_model(model_url, num_classes=10):
...
feature_extractor_layer = ...

# Create our own model
  if num_classes > 2:
    out_put_layer = layers.Dense(num_classes, activation='softmax', name='output_layer')
  else:
    out_put_layer = layers.Dense(1, activation='sigmoid', name='output_layer')
  
  model = tf.keras.Sequential([
      feature_extractor_layer,
      out_put_layer
  ])