tensorflow2使用预训练模型和权重进行图像分类

使用回调函数保存模型和自定义训练终止。

数据集：

 

import warnings
warnings.filterwarnings("ignore")
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import keras.layers
from tensorflow.keras import (models,layers,callbacks,utils,)
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
import pathlib
from PIL import Image



def load_dataset():
    """ 载入数据\n
        """
    data_dir = "flower_photos"
    data_dir = pathlib.Path(data_dir)

    image_count = len(list(data_dir.glob('*/*.jpg')))
    print('total images number:',image_count)

    roses = list(data_dir.glob('roses/*'))
    aimage = Image.open(str(roses[1]))
    # aimage.show()

    batch_size = 32


    train_ds = utils.image_dataset_from_directory(data_dir,
                                                  validation_split=0.2,
                                                  subset="training",
                                                  seed=123,
                                                  image_size=(img_height, img_width),
                                                  batch_size=batch_size)
    val_ds = utils.image_dataset_from_directory(data_dir,
                                                validation_split=0.2,
                                                subset="validation",
                                                seed=123,
                                                image_size=(img_height, img_width),
                                                batch_size=batch_size)

    class_names = train_ds.class_names
    print('class names:',class_names)

    return train_ds,val_ds,class_names

def show_some_images():
    plt.figure()
    for images, labels in train_ds.take(2):
        for i in range(16):
            ax = plt.subplot(4,4,i+1)
            plt.imshow(images[i].numpy().astype('uint8'))
            plt.title(classnames[labels[i]])
            plt.axis('off')
        plt.show()

def create_a_sequenceModel():
    """ 建立一个模型 """
    num_classes = 5
    # model = models.Sequential([
    #     # tf.keras.layers.Rescaling(1./255, input_shape=(img_height,img_width,3)),
    #     layers.Conv2D(16,3,padding='same',activation='relu',input_shape=(img_height,img_width,3)),
    #     layers.MaxPooling2D(),
    #     layers.Dropout(0.2),
    #     layers.Conv2D(32,3,padding='same',activation='relu'),
    #     layers.MaxPooling2D(),
    #     layers.Dropout(0.2),
    #     layers.Conv2D(64,3,padding='same',activation='relu'),
    #     layers.MaxPooling2D(),
    #     layers.Flatten(),
    #     layers.Dense(128,activation='relu'),
    #     layers.Dense(num_classes),
    # ])

    #todo: create a model from keras
    # model = tf.keras.applications.densenet.DenseNet201(input_shape=(img_height,img_width,3),
    #                                           include_top=True,
    #                                           weights=None,
    #                                           classes=5)
    ##todo: 使用预训练权重
    base_model = tf.keras.applications.densenet.DenseNet201(weights='imagenet',
                                                            include_top=False)
    x = base_model.output
    x = keras.layers.GlobalAveragePooling2D()(x)
    x = keras.layers.Dense(1024,activation='relu')(x)
    predictions = keras.layers.Dense(num_classes,activation='softmax')(x)
    model = keras.models.Model(inputs=base_model.input,outputs=predictions)

    return model

class stopCallback(tf.keras.callbacks.Callback):
    """ stop callback class """
    def on_epoch_end(self, epoch, logs=None):
        if logs is None:
            logs = {}
        if(logs.get('loss')<0.4):
          print("\nReached 60% accuracy so cancelling training!")
          self.model.stop_training = True


if __name__ == '__main__':
    img_height = 180
    img_width = 180
    train_ds, val_ds, classnames = load_dataset()
    # show_some_images()
    # for image_batch, labels_batch in train_ds:
    #     print(image_batch.shape)
    #     print(labels_batch.shape)
    #     break

    AUTOTUNE = tf.data.AUTOTUNE
    train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
    val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

    #todo: standardize the data
    normalization_layer = tf.keras.layers.Rescaling(1./255)
    normalized_ds = train_ds.map(lambda x,y:(normalization_layer(x),y))
    image_batch, labels_batch = next(iter(normalized_ds))
    first_image = image_batch[0]
    print('min and max:',np.min(first_image),np.max(first_image))

    model = create_a_sequenceModel()


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

    #todo: set save callbacks
    logdir = 'logs/'
    if not os.path.exists(logdir):
        os.mkdir(logdir)
    checkpoint_path = logdir+"EP{epoch:03d}_acc{accuracy:.3f}-vac{val_accuracy:.3f}.h5"
    checkpoint_dir = os.path.dirname(checkpoint_path)
    save_callback = callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                            save_best_only=True,
                                            save_weights_only=True,
                                              verbose=1)
    #todo: set a stop callback
    astop_callback = stopCallback()


    #todo: train the model
    epochs = 10
    history = model.fit(train_ds,
                        validation_data=val_ds,
                        epochs=epochs,
                        callbacks=[save_callback,astop_callback])