导入所需库

import os import random import numpy as np import pandas as pd import tensorflow as tf import matplotlib.pyplot as plt from tensorflow.keras.models import Sequential, Model from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D, BatchNormalization, GlobalAveragePooling2D from tensorflow.keras.preprocessing.image import ImageDataGenerator from tensorflow.keras.preprocessing import image from tensorflow.keras.applications.vgg16 import VGG16

设置训练集和测试集目录

train_dir = 'C:/Users/chaofan/Desktop/class/class/archive/seg_train/seg_train' # 训练集目录 val_dir = 'C:/Users/chaofan/Desktop/class/class/archive/seg_test/seg_test' # 测试集目录

获取类别列表

classes = os.listdir(train_dir)

设置批大小和图像尺寸

batch_size = 64 # 批大小 IMG_HEIGHT = 150 # 图像高度 IMG_WIDTH = 150 # 图像宽度

设置训练轮数

epochs = 10 # 训练轮数

创建训练集和测试集的图像数据生成器

train_image_generator = ImageDataGenerator( rescale=1./255, # 归一化 horizontal_flip=True # 水平翻转 )

val_image_generator = ImageDataGenerator( rescale=1./255 # 归一化 )

从目录中加载训练集和测试集数据

train_data_gen = train_image_generator.flow_from_directory( batch_size=batch_size, # 批大小 directory=train_dir, # 训练集目录 shuffle=True, # 是否打乱数据 target_size=(IMG_HEIGHT, IMG_WIDTH), # 图像大小 class_mode='categorical' # 分类方式 )

val_data_gen = val_image_generator.flow_from_directory( batch_size=batch_size, # 批大小 directory=val_dir, # 验证集目录 shuffle=True, # 是否打乱数据 target_size=(IMG_HEIGHT, IMG_WIDTH), # 图像大小 class_mode='categorical' # 分类方式 )

获取训练集和验证集的样本数

total_train = len(train_data_gen) # 训练集样本数 total_val = len(val_data_gen) # 验证集样本数

print('总训练数据批次数：', total_train) print('总验证数据批次数： ', total_val)

使用预训练的VGG16模型作为基础模型

base_model = VGG16(weights='imagenet', include_top=False, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3))

冻结基础模型的所有层

for layer in base_model.layers: layer.trainable = False

创建顺序模型

model = Sequential()

添加基础模型

model.add(base_model)

添加扁平层

model.add(Flatten())

添加全连接层

model.add(Dense(1024, activation='relu')) model.add(Dense(512, activation='relu')) model.add(Dense(64, activation='relu'))

添加输出层

model.add(Dense(len(classes), activation='softmax'))

编译模型

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

打印模型结构

model.summary()

训练模型

history = model.fit( train_data_gen, # 训练集数据生成器 steps_per_epoch=total_train, # 每轮迭代的步数 epochs=epochs, # 训练轮数 validation_data=val_data_gen, # 验证集数据生成器 validation_steps=total_val # 验证集每轮迭代的步数 )

绘制训练过程中的准确率和损失曲线

acc = history.history['accuracy'] val_acc = history.history['val_accuracy'] loss = history.history['loss'] val_loss = history.history['val_loss']

plt.figure(figsize=(8, 8)) plt.subplot(2, 1, 1) plt.plot(acc, label='Training Accuracy') plt.plot(val_acc, label='Validation Accuracy') plt.legend(loc='lower right') plt.ylabel('Accuracy') plt.ylim([min(plt.ylim()), 1]) plt.title('Training and Validation Accuracy')

plt.subplot(2, 1, 2) plt.plot(loss, label='Training Loss') plt.plot(val_loss, label='Validation Loss') plt.legend(loc='upper right') plt.ylabel('Cross Entropy') plt.ylim([0, 1.0]) plt.title('Training and Validation Loss') plt.xlabel('epoch')

plt.show()