项目概述

本项目展示了如何使用预训练的 EfficientNet 模型在斯坦福狗数据集上进行图像分类任务。
基于 Tan 和 Le, 2019 提出的 EfficientNet 架构，
提供了完整的迁移学习工作流程，从数据准备到模型训练和评估。

什么是 EfficientNet

根据 main.py 中的介绍：

EfficientNet 由 Tan 和 Le, 2019 首次提出，
是目前最高效的模型之一（即在推理时需要最少的 FLOPS 计算量），
在 ImageNet 和常见的图像分类迁移学习任务上都达到了最先进的准确率。

EfficientNet 提供了一系列模型变体（B0 到 B7），这些模型在各种规模上都代表了效率和准确性的良好结合。

EfficientNet 的 B0 到 B7 变体

下表列出了不同 EfficientNet 变体的输入分辨率要求：

基础模型	分辨率
EfficientNetB0	224
EfficientNetB1	240
EfficientNetB2	260
EfficientNetB3	300
EfficientNetB4	380
EfficientNetB5	456
EfficientNetB6	528
EfficientNetB7	600

安装指南

环境要求

• Python 3.7 或更高版本
• TensorFlow 2.x
• Keras 2.3 或更高版本
• CUDA 支持（可选，用于 GPU 加速）

安装依赖

pip install numpy tensorflow tensorflow-datasets matplotlib

验证安装

python -c "import tensorflow as tf; print(tf.__version__)"
python -c "import tensorflow_datasets as tfds; print('TFDS 已安装')"

依赖要求

项目依赖以下库：

import numpy as np
import tensorflow_datasets as tfds
import tensorflow as tf
import matplotlib.pyplot as plt
import keras
from keras import layers
from keras.applications import EfficientNetB0

数据集

本项目使用 斯坦福狗 数据集，包含属于 120 个犬种的 20,580 张图像：

• 训练集：12,000 张图像
• 测试集：8,580 张图像

数据集加载代码：

dataset_name = "stanford_dogs"
(ds_train, ds_test), ds_info = tfds.load(
    dataset_name, split=["train", "test"], with_info=True, as_supervised=True
)
NUM_CLASSES = ds_info.features["label"].num_classes

使用方法

1. 数据预处理

图像预处理和增强流程如下：

# 图像大小设置
IMG_SIZE = 224
BATCH_SIZE = 64

# 调整图像大小
size = (IMG_SIZE, IMG_SIZE)
ds_train = ds_train.map(lambda image, label: (tf.image.resize(image, size), label))
ds_test = ds_test.map(lambda image, label: (tf.image.resize(image, size), label))

# 数据增强层
img_augmentation_layers = [
    layers.RandomRotation(factor=0.15),
    layers.RandomTranslation(height_factor=0.1, width_factor=0.1),
    layers.RandomFlip(),
    layers.RandomContrast(factor=0.1),
]

def img_augmentation(images):
    for layer in img_augmentation_layers:
        images = layer(images)
    return images

2. 模型构建

2.1 从头开始训练

如果要从零开始训练模型（不推荐，除非有充足的计算资源）：

model = EfficientNetB0(
    include_top=True,
    weights=None,
    classes=NUM_CLASSES,
    input_shape=(IMG_SIZE, IMG_SIZE, 3),
)
model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])

2.2 迁移学习

推荐使用预训练权重进行迁移学习，这可以显著提高性能并减少训练时间：

def build_model(num_classes):
    inputs = layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
    model = EfficientNetB0(include_top=False, input_tensor=inputs, weights="imagenet")

    # 冻结预训练权重
    model.trainable = False

    # 重建顶层
    x = layers.GlobalAveragePooling2D(name="avg_pool")(model.output)
    x = layers.BatchNormalization()(x)

    top_dropout_rate = 0.2
    x = layers.Dropout(top_dropout_rate, name="top_dropout")(x)
    outputs = layers.Dense(num_classes, activation="softmax", name="pred")(x)

    # 编译
    model = keras.Model(inputs, outputs, name="EfficientNet")
    optimizer = keras.optimizers.Adam(learning_rate=1e-2)
    model.compile(
        optimizer=optimizer, loss="categorical_crossentropy", metrics=["accuracy"]
    )
    return model

3. 微调模型

在训练顶层后，可以解冻部分层进行微调以进一步提高性能：

def unfreeze_model(model):
    # 解冻顶部 20 层，同时保持 BatchNorm 层冻结
    for layer in model.layers[-20:]:
        if not isinstance(layer, layers.BatchNormalization):
            layer.trainable = True

    optimizer = keras.optimizers.Adam(learning_rate=1e-5)
    model.compile(
        optimizer=optimizer, loss="categorical_crossentropy", metrics=["accuracy"]
    )

训练流程

本项目采用两阶段训练策略：

阶段 1：冻结预训练层

• 冻结特征提取器，仅训练顶层
• 使用较大学习率 (1e-2)
• 训练约 25 轮

阶段 2：微调

• 解冻部分层进行微调（保持 BatchNorm 层冻结）
• 使用较小学习率 (1e-5)
• 训练约 4 轮

微调技巧

根据 main.py 中的注释，微调 EfficientNet 时的关键技巧：

层冻结策略

• BatchNormalization 层需要保持冻结：如果它们也被设为可训练，解冻后的第一个轮次会显著降低准确率。
• 选择性解冻：在某些情况下，只开放部分层而不是解冻所有层可能更有益。
• 块级解冻：每个块需要全部开启或关闭。这是因为架构包括每个块从第一层到最后一层的快捷连接。

超参数选择

• 模型选择：EfficientNet 的更大变体并不保证性能改善，特别是对于数据较少或类别较少的任务。
• 优化器选择：不要在迁移学习中使用原始论文中的 RMSprop 设置，对于迁移学习，动量和学习率太高。
• 批量大小：较小的批量大小有利于验证准确率，可能是因为有效地提供了正则化。

运行说明

快速开始

1. 克隆或下载本项目到本地
2. 安装所需依赖
3. 运行主脚本：

python main.py

训练流程说明

运行 main.py 将会执行以下步骤：

1. 数据加载和预处理：自动下载斯坦福狗数据集并进行预处理
2. 数据增强：应用随机旋转、平移、翻转和对比度调整等增强技术
3. 模型训练：首先冻结预训练模型的特征提取层，仅训练顶层
4. 模型微调：解冻部分层进行微调以进一步提高性能
5. 结果可视化：绘制训练和验证准确率曲线

自定义训练参数

您可以在 main.py 中修改以下参数来调整训练过程：

# 图像大小设置
IMG_SIZE = 224  # 根据选择的 EfficientNet 变体调整
BATCH_SIZE = 64  # 根据可用内存调整

# 训练轮次
epochs = 25  # 冻结训练阶段
epochs = 4   # 微调阶段

使用不同的数据集

您可以通过修改 dataset_name 参数来使用 TFDS 中的其他数据集：

dataset_name = "stanford_dogs"  # 可替换为 "cifar10", "cifar100", "food101" 等

使用不同的 EfficientNet 变体

如需使用更大或更小的 EfficientNet 变体，请修改模型导入和创建部分：

# 例如，使用 EfficientNetB3
from keras.applications import EfficientNetB3
IMG_SIZE = 300  # B3 变体对应的输入大小

# 在 build_model 函数中替换
model = EfficientNetB3(include_top=False, input_tensor=inputs, weights="imagenet")

可视化训练结果

使用以下函数绘制训练和验证准确率：

def plot_hist(hist):
    plt.plot(hist.history["accuracy"])
    plt.plot(hist.history["val_accuracy"])
    plt.title("模型准确率")
    plt.ylabel("准确率")
    plt.xlabel("轮次")
    plt.legend(["训练", "验证"], loc="upper left")
    plt.show()

完整代码

# 导入必要的库
import numpy as np                  # 用于数值计算和数组操作
import tensorflow_datasets as tfds  # 用于加载预构建数据集
import tensorflow as tf             # 深度学习框架
import matplotlib.pyplot as plt     # 用于绘制图表和可视化
import keras                        # 高级神经网络 API
from keras import layers            # 神经网络层
from keras.applications import EfficientNetB0  # 预训练的 EfficientNetB0 模型


## 数据加载和预处理

#本节定义了数据处理的核心参数和初始数据加载流程。
#我们使用斯坦福狗数据集，这是一个包含120个狗品种的图像分类数据集。


# 定义关键参数
IMG_SIZE = 224      # 图像尺寸，EfficientNetB0要求的输入大小
BATCH_SIZE = 64     # 批次大小，控制训练时的内存使用和训练速度
dataset_name = "stanford_dogs"  # 使用的数据集名称

# 从TensorFlow Datasets加载斯坦福狗数据集
# split参数指定数据集的分割方式，"train"和"test"分别表示训练集和测试集
# with_info=True返回数据集的元信息
# as_supervised=True返回(image, label)的元组对
(ds_train, ds_test), ds_info = tfds.load(
    dataset_name, split=["train", "test"], with_info=True, as_supervised=True
)

# 获取数据集的类别数量
NUM_CLASSES = ds_info.features["label"].num_classes
print(f"数据集包含 {NUM_CLASSES} 个类别（狗品种）")

# 定义调整大小的尺寸
size = (IMG_SIZE, IMG_SIZE)

# 调整所有训练集图像的大小
# 使用map函数对每个图像-标签对应用resize操作
# tf.image.resize将图像调整为指定的尺寸，保持宽高比，使用双线性插值
print("正在调整训练集图像大小...")
ds_train = ds_train.map(lambda image, label: (tf.image.resize(image, size), label))

# 调整所有测试集图像的大小
print("正在调整测试集图像大小...")
ds_test = ds_test.map(lambda image, label: (tf.image.resize(image, size), label))


## 数据可视化

#在处理数据之前，先可视化一些样本图像，以了解数据集的基本情况。

# 定义标签格式化函数，将标签ID转换为易读的品种名称
def format_label(label):
    # 将整数标签转换为字符串形式
    string_label = label_info.int2str(label)
    # 分割字符串并提取品种名称部分（移除前缀数字）
    return string_label.split("-")[1]

# 获取标签信息，用于将整数标签转换为字符串形式
label_info = ds_info.features["label"]

# 可视化训练集中的前9个图像
print("正在可视化训练集样本图像...")
plt.figure(figsize=(10, 10))  # 设置图表大小
for i, (image, label) in enumerate(ds_train.take(9)):  # 只取前9个样本
    ax = plt.subplot(3, 3, i + 1)  # 创建3x3的子图网格
    plt.imshow(image.numpy().astype("uint8"))  # 显示图像
    plt.title("{}".format(format_label(label)))  # 设置图像标题为品种名称
    plt.axis("off")  # 不显示坐标轴
plt.tight_layout()  # 调整布局
plt.show()  # 显示图表


## 数据增强

#数据增强是提高模型泛化能力的重要技术，通过对训练图像进行随机变换，
#可以有效增加数据集的多样性，减少过拟合。


# 定义数据增强层列表
# 这些层将用于随机变换训练图像
img_augmentation_layers = [
    layers.RandomRotation(factor=0.15),           # 随机旋转，角度范围[-15%, 15%]
    layers.RandomTranslation(height_factor=0.1, width_factor=0.1),  # 随机平移
    layers.RandomFlip(),                         # 随机水平翻转
    layers.RandomContrast(factor=0.1),           # 随机对比度调整
]

# 定义图像增强函数
def img_augmentation(images):
    # 对输入图像依次应用所有增强变换
    for layer in img_augmentation_layers:
        images = layer(images)
    return images

# 可视化数据增强效果
print("正在可视化数据增强效果...")
plt.figure(figsize=(10, 10))  # 设置图表大小
for image, label in ds_train.take(1):  # 只取一个样本进行增强可视化
    for i in range(9):  # 生成9个不同的增强版本
        ax = plt.subplot(3, 3, i + 1)
        # 扩展维度以适应增强函数的输入要求
        aug_img = img_augmentation(np.expand_dims(image.numpy(), axis=0))
        aug_img = np.array(aug_img)  # 转换为NumPy数组
        plt.imshow(aug_img[0].astype("uint8"))  # 显示增强后的图像
        plt.title("{}".format(format_label(label)))  # 设置标题
        plt.axis("off")  # 不显示坐标轴
plt.tight_layout()  # 调整布局
plt.show()  # 显示图表


## 数据预处理和批处理

#定义训练集和测试集的预处理函数，并将数据集转换为批处理形式，
#以提高训练效率。


# 定义训练集预处理函数
def input_preprocess_train(image, label):
    # 应用数据增强
    image = img_augmentation(image)
    # 将标签转换为独热编码格式
    # 独热编码是分类任务中表示类别标签的常用方式
    label = tf.one_hot(label, NUM_CLASSES)
    return image, label

# 定义测试集预处理函数（不需要数据增强）
def input_preprocess_test(image, label):
    # 只需要将标签转换为独热编码
    label = tf.one_hot(label, NUM_CLASSES)
    return image, label

# 应用预处理函数到训练集
# num_parallel_calls=tf.data.AUTOTUNE 自动确定并行处理的线程数
print("正在预处理训练集...")
ds_train = ds_train.map(input_preprocess_train, num_parallel_calls=tf.data.AUTOTUNE)
# 将训练集分成批次，每批次BATCH_SIZE个样本
# drop_remainder=True 丢弃最后不足一个批次的样本，确保批次大小一致
print("正在将训练集分成批次...")
ds_train = ds_train.batch(batch_size=BATCH_SIZE, drop_remainder=True)
# 预取数据以提高训练效率，tf.data.AUTOTUNE自动调整预取量
print("正在配置训练集预取...")
ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

# 应用预处理函数到测试集
print("正在预处理测试集...")
ds_test = ds_test.map(input_preprocess_test, num_parallel_calls=tf.data.AUTOTUNE)
# 将测试集分成批次
print("正在将测试集分成批次...")
ds_test = ds_test.batch(batch_size=BATCH_SIZE, drop_remainder=True)

model = EfficientNetB0(
    include_top=True,
    weights=None,
    classes=NUM_CLASSES,
    input_shape=(IMG_SIZE, IMG_SIZE, 3),
)
model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])
model.summary()

epochs = 40
hist = model.fit(ds_train, epochs=epochs, validation_data=ds_test)

def plot_hist(hist):
    plt.plot(hist.history["accuracy"])
    plt.plot(hist.history["val_accuracy"])
    plt.title("模型准确率")
    plt.ylabel("准确率")
    plt.xlabel("轮次")
    plt.legend(["训练", "验证"], loc="upper left")
    plt.show()

plot_hist(hist)

## 模型构建

#本节定义了如何使用迁移学习构建分类模型。
#我们使用预训练的EfficientNetB0作为基础模型，
#并在其基础上添加一个分类头部，实现对狗品种的识别。

# 定义模型构建函数
def build_model(num_classes):
    # 创建输入层，指定输入图像的形状
    # EfficientNetB0要求输入为224x224，RGB三通道
    inputs = layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
    
    # 加载预训练的EfficientNetB0模型作为特征提取器
    # include_top=False 表示不包含顶层分类器
    # input_tensor=inputs 指定输入张量
    # weights="imagenet" 表示使用在ImageNet数据集上预训练的权重
    model = EfficientNetB0(include_top=False, input_tensor=inputs, weights="imagenet")
    
    # 冻结基础模型的权重，避免在初始训练时破坏预训练的特征提取能力
    model.trainable = False
    
    # 添加全局平均池化层，将特征图转换为特征向量
    # 这比全连接层更高效，能减少过拟合
    x = layers.GlobalAveragePooling2D(name="avg_pool")(model.output)
    
    # 添加批归一化层，稳定训练过程，提高收敛速度
    x = layers.BatchNormalization()(x)
    
    # 添加丢弃层，丢弃比例为0.2，减少过拟合
    top_dropout_rate = 0.2
    x = layers.Dropout(top_dropout_rate, name="top_dropout")(x)
    
    # 添加最终的分类层，使用softmax激活函数输出概率分布
    outputs = layers.Dense(num_classes, activation="softmax", name="pred")(x)
    
    # 创建完整的模型，指定输入和输出
    model = keras.Model(inputs, outputs, name="EfficientNet")
    
    # 配置优化器，使用Adam优化器，学习率设为0.01
    optimizer = keras.optimizers.Adam(learning_rate=1e-2)
    
    # 编译模型，配置损失函数和评估指标
    model.compile(optimizer=optimizer, loss="categorical_crossentropy", metrics=["accuracy"])
    
    return model

# 构建模型，传入类别数量
print("正在构建模型...")
model = build_model(num_classes=NUM_CLASSES)

## 第一阶段：训练分类头部

#在第一阶段，我们冻结了预训练模型的权重，
#只训练我们添加的分类头部。

# 定义训练轮数
epochs = 25

# 开始训练模型
print("开始第一阶段训练（只训练分类头部）...")
hist = model.fit(
    ds_train,           # 训练数据集
    epochs=epochs,      # 训练轮数
    validation_data=ds_test  # 验证数据集，用于在每轮结束时评估模型性能
)

# 绘制训练过程中的准确率变化
plot_hist(hist)

## 第二阶段：微调模型

#在第二阶段，我们解冻预训练模型的最后20层权重，
#使用较小的学习率对整个模型进行微调，以适应特定任务。

# 定义模型解冻和重新编译函数
def unfreeze_model(model):
    # 解冻模型中的最后20层，除了批归一化层
    # 保持批归一化层冻结是一种常用的微调技巧，可以避免破坏其统计信息
    # 只微调最后几层可以节省计算资源，同时允许模型适应特定任务
    for layer in model.layers[-20:]:
        if not isinstance(layer, layers.BatchNormalization):
            layer.trainable = True
    
    # 使用较小的学习率，避免破坏预训练权重
    # 微调阶段通常使用比初始训练更小的学习率
    optimizer = keras.optimizers.Adam(learning_rate=1e-5)
    
    # 重新编译模型
    model.compile(optimizer=optimizer, loss="categorical_crossentropy", metrics=["accuracy"])

# 解冻模型并重新编译
print("正在解冻模型并重新编译...")
unfreeze_model(model)

# 定义微调的轮数
# 微调通常需要较少的轮次
print("开始第二阶段微调（训练模型最后20层）...")
epochs = 4

# 开始微调模型
hist = model.fit(
    ds_train,           # 训练数据集
    epochs=epochs,      # 微调轮数
    validation_data=ds_test  # 验证数据集
)

# 绘制微调过程中的准确率变化
plot_hist(hist)