使用 VGG16 进行 MNIST 数字识别的迁移学习

碧海醫心

发布时间：2025-08-19 19:12:16

719人浏览过

来源于php中文网

原创

使用 vgg16 进行 mnist 数字识别的迁移学习

本文档旨在指导读者如何利用 VGG16 模型进行 MNIST 手写数字识别的迁移学习。我们将重点介绍如何构建模型，加载预训练权重，以及解决可能出现的 GPU 配置问题。本文将帮助您理解迁移学习的基本原理，并提供一个可运行的示例，用于 MNIST 数据集的分类。

使用 VGG16 进行 MNIST 数字识别的迁移学习

迁移学习是一种强大的技术，它允许我们将一个在大型数据集上训练过的模型（例如，ImageNet）的知识迁移到另一个相关但更小的数据集上（例如，MNIST）。这种方法可以显著提高模型的性能，尤其是在数据量有限的情况下。本文将介绍如何使用 VGG16 模型进行 MNIST 手写数字识别的迁移学习。

准备工作

在开始之前，请确保您已安装以下库：

TensorFlow
Keras
NumPy

import tensorflow as tf
from tensorflow.keras.applications import VGG16
from tensorflow.keras import layers, models
import numpy as np

构建迁移学习模型

以下代码展示了如何使用 VGG16 模型作为基础模型，并添加自定义层以适应 MNIST 数据集。

class VGG16TransferLearning(tf.keras.Model):
    def __init__(self, base_model, models):
        super(VGG16TransferLearning, self).__init__()
        # base model
        self.base_model = base_model

        # other layers
        self.flatten = tf.keras.layers.Flatten()
        self.dense1 = tf.keras.layers.Dense(512, activation='relu')
        self.dense2 = tf.keras.layers.Dense(512, activation='relu')
        self.dense3 = tf.keras.layers.Dense(10)
        self.layers_list = [self.flatten, self.dense1, self.dense2, self.dense3]

        # instantiate the base model with other layers
        self.model = models.Sequential(
            [self.base_model, *self.layers_list]
        )

    def call(self, *args, **kwargs):
        out = args[0]
        for layer in self.model.layers:
            out = layer(out)
        return out

在这个类中，VGG16TransferLearning 继承自 tf.keras.Model。__init__ 方法初始化了 VGG16 基础模型，以及一些自定义的 Dense 层。call 方法定义了模型的前向传播过程。

加载 VGG16 预训练权重

在实例化 VGG16TransferLearning 类之前，我们需要加载 VGG16 模型并冻结其权重，以防止在训练过程中修改预训练的权重。

# 假设 x_train[0].shape 是 (75, 75, 3)
input_shape = (75, 75, 3)  # 确保输入形状与您的数据一致

base_model = VGG16(weights="imagenet", include_top=False, input_shape=input_shape)
base_model.trainable = False  # 冻结 VGG16 的权重

VGG16(weights="imagenet", include_top=False, input_shape=x_train[0].shape) 这行代码加载了在 ImageNet 上预训练的 VGG16 模型。include_top=False 表示我们不包含 VGG16 的顶层分类器，因为我们需要根据 MNIST 数据集进行自定义。 base_model.trainable = False 冻结了 VGG16 模型的权重，防止在训练过程中更新这些权重。

编译和训练模型

接下来，我们将实例化 VGG16TransferLearning 类，编译模型，并使用 MNIST 数据集进行训练。

model = VGG16TransferLearning(base_model, models)

model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              optimizer=tf.keras.optimizers.legacy.Adam(),
              metrics=['accuracy'])

# 假设 x_train, y_train, x_test, y_test 已经加载
# 并且 x_train 的形状是 (num_samples, 75, 75, 3)
# 并且 y_train 是整数标签
model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

model.compile 方法配置了模型的训练过程，包括损失函数、优化器和评估指标。model.fit 方法使用训练数据训练模型，并使用验证数据评估模型的性能。

GPU 配置问题

在运行上述代码时，可能会遇到 GPU 配置问题，导致 Kernel 死机。这通常是因为 TensorFlow 没有正确检测到 GPU。以下是一些可能的解决方案：

检查 TensorFlow 是否检测到 GPU：

PHP高级开发技巧与范例
PHP是一种功能强大的网络程序设计语言，而且易学易用，移植性和可扩展性也都非常优秀，本书将为读者详细介绍PHP编程。全书分为预备篇、开始篇和加速篇三大部分，共9章。预备篇主要介绍一些学习PHP语言的预备知识以及PHP运行平台的架设；开始篇则较为详细地向读者介绍PKP语言的基本语法和常用函数，以及用PHP如何对MySQL数据库进行操作；加速篇则通过对典型实例的介绍来使读者全面掌握PHP。本书

下载
```
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
```
如果输出为 0，则表示 TensorFlow 没有检测到 GPU。
安装正确的 TensorFlow 版本：

确保您安装了与您的 GPU 兼容的 TensorFlow 版本。例如，如果您使用的是 NVIDIA GPU，则需要安装 tensorflow-gpu 版本。
配置 CUDA 和 cuDNN：

确保您已正确安装和配置 CUDA 和 cuDNN，并且它们的版本与 TensorFlow 兼容。
设置环境变量：

设置以下环境变量：
- CUDA_HOME: CUDA 的安装路径。
- LD_LIBRARY_PATH: CUDA 库的路径。