shark-classifier / models /model_v1.py
zihaoz96's picture
first commit
961927f
raw
history blame
1.72 kB
import tensorflow as tf
from tensorflow.keras.callbacks import TensorBoard, EarlyStopping, ModelCheckpoint
from tensorflow.keras.layers import Conv2D, Dense, GlobalMaxPooling2D
from tensorflow.keras.layers import Dense, MaxPooling2D, BatchNormalization
from tensorflow.keras.models import Sequential
from tensorflow.keras import Model
def model_v1(nbr_class):
model = Sequential()
model.add(Conv2D(64,(3,3), activation="relu", input_shape=(224,224,3)))
model.add(BatchNormalization())
model.add(Conv2D(64,(3,3), activation="relu"))
model.add(BatchNormalization())
model.add(MaxPooling2D())
model.add(Conv2D(128,(3,3), activation="relu"))
model.add(BatchNormalization())
model.add(Conv2D(128,(3,3), activation="relu"))
model.add(BatchNormalization())
model.add(MaxPooling2D())
model.add(Conv2D(256,(3,3), activation="relu"))
model.add(BatchNormalization())
model.add(Conv2D(256,(3,3), activation="relu"))
model.add(BatchNormalization())
model.add(MaxPooling2D())
# model.add(Conv2D(512,(3,3), activation="relu"))
# model.add(BatchNormalization())
# model.add(Conv2D(512,(3,3), activation="relu"))
# model.add(BatchNormalization())
# model.add(MaxPooling2D())
# model.add(Conv2D(512,(3,3), activation="relu"))
# model.add(BatchNormalization())
# model.add(Conv2D(512,(3,3), activation="relu"))
# model.add(BatchNormalization())
# model.add(Conv2D(512,(3,3), activation="relu"))
# model.add(BatchNormalization())
# model.add(GlobalMaxPooling2D())
model.add(Dense(1024, activation="relu"))
model.add(BatchNormalization())
model.add(Dense(nbr_class, activation="softmax"))
return model