import os
import warnings
warnings.filterwarnings("ignore")
import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.applications.resnet import ResNet50

#数据所在文件夹
base_dir = './data/cats_and_dogs'
train_dir = os.path.join(base_dir,'train')

train_data,validation_data = tf.keras.utils.image_dataset_from_directory(
    train_dir,
    labels='inferred',
    label_mode='binary',
    class_names=["cats","dogs"],
    color_mode='rgb',
    batch_size=4,
    image_size=(64, 64),
    shuffle=True,
    seed=2023,
    validation_split=0.5,
    subset='both',
    interpolation='bilinear',
    follow_links=False,
    crop_to_aspect_ratio=False,
)

save_model_cb = tf.keras.callbacks.ModelCheckpoint(filepath='model_resnet50_cats_and_dogs.h5', save_freq='epoch')

base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(64, 64, 3))
base_model.trainable = True
    
model = tf.keras.models.Sequential([
    base_model,
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation='relu',kernel_regularizer=tf.keras.regularizers.l2(l=0.01)),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

model.compile(loss='binary_crossentropy',optimizer = Adam(lr=1e-3),metrics = ['acc'])

model.summary()
history = model.fit(train_data.repeat(),steps_per_epoch=100,epochs=50,validation_data=validation_data.repeat(),validation_steps=50,verbose=1,callbacks = [save_model_cb])