Skip to content

Instantly share code, notes, and snippets.

@alonlavian
Last active May 28, 2019 13:22
Show Gist options
  • Save alonlavian/cf4c6e85aa7f042f15f09da65fc14182 to your computer and use it in GitHub Desktop.
Save alonlavian/cf4c6e85aa7f042f15f09da65fc14182 to your computer and use it in GitHub Desktop.
from tensorflow.keras.applications.vgg16 import VGG16
from tensorflow.keras import layers
from tensorflow.keras.models import Model
base_model = VGG16(weights='imagenet', #pre-trained weights
include_top=False, #remove original classifier
input_shape=(HEIGHT, WIDTH, 3))
#Freeze all layers
for layer in base_model.layers:
layer.trainable = False
#Add binary classifier
x = base_model.output
x = layers.Flatten()(base_model.output)
x = layers.Dense(1024, activation='relu')(x)
x = layers.Dropout(0.2)(x)
x = layers.Dense (1, activation='sigmoid')(x)
model = Model(inputs=base_model.input, outputs=x)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment