基于VGG16的MNIST手写数字迁移学习教程

聖光之護

发布时间：2025-08-19 19:42:00

903人浏览过

来源于php中文网

原创

基于vgg16的mnist手写数字迁移学习教程

本文将介绍如何使用VGG16模型进行MNIST手写数字的迁移学习，重点在于解决在训练过程中可能遇到的Kernel崩溃问题。我们将详细分析代码结构，并提供一种可能的解决方案，即确保TensorFlow正确识别并使用GPU资源，从而顺利完成模型的训练和验证，并为后续基于梯度的攻击提供logits。

迁移学习与VGG16

迁移学习是一种机器学习方法，它允许我们将从一个问题上获得的知识应用到另一个相关问题上。在图像识别领域，预训练的卷积神经网络（CNN），如VGG16，已经在ImageNet等大型数据集上学习到了丰富的特征表示。我们可以利用这些预训练的权重，将其应用到新的图像识别任务中，从而加快训练速度并提高模型性能。

代码结构分析

以下代码实现了一个基于VGG16的迁移学习模型，用于MNIST手写数字识别：

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

class VGG16TransferLearning(tf.keras.Model):
  def __init__(self, base_model, models):
    super(VGG16TransferLearning, self).__init__()
    # base model
    self.base_model = base_model
    self.base_model.trainable = False  # 冻结预训练模型的权重

    # other layers
    self.flatten = layers.Flatten()
    self.dense1 = layers.Dense(512, activation='relu')
    self.dense2 = layers.Dense(512, activation='relu')
    self.dense3 = 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, inputs, training=None, mask=None):
    x = inputs
    for layer in self.model.layers:
        x = layer(x)
    if training:
        return x
    else:
        prob = tf.nn.softmax(x)
        return x, prob

# 加载预训练的VGG16模型
base_model = VGG16(weights="imagenet", include_top=False, input_shape=(75, 75, 3))

# 构建模型
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 已经加载)
# 确保输入数据的形状与 base_model.input_shape 匹配

# 训练模型
# model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

代码解释：

VGG16TransferLearning 类: 继承自 tf.keras.Model，用于构建迁移学习模型。
- __init__ 方法：
  - 接收一个预训练的 base_model (VGG16) 作为输入。
  - 冻结 base_model 的权重，防止在训练过程中被修改。这是迁移学习的关键步骤，保留了预训练模型学习到的通用特征。
  - 定义了新的全连接层 (flatten, dense1, dense2, dense3)，用于适应MNIST数据集的分类任务。
  - 使用 tf.keras.models.Sequential 将 base_model 和新的全连接层组合成完整的模型。
- call 方法：定义了模型的前向传播过程。
  - 遍历模型的每一层，将输入 x 依次通过每一层。
  - 根据 training 参数，返回 logits (训练时) 或 logits 和 softmax 概率 (推理时)。
加载VGG16模型: 使用 VGG16(weights="imagenet", include_top=False, input_shape=(75, 75, 3)) 加载预训练的 VGG16 模型。
- weights="imagenet"：指定使用在 ImageNet 数据集上预训练的权重。
- include_top=False：移除 VGG16 模型的顶层全连接层，因为我们需要自定义全连接层来适应 MNIST 数据集。
- input_shape=(75, 75, 3)：指定输入图像的形状。MNIST 数据集是灰度图像，但 VGG16 期望输入 RGB 图像，因此需要将灰度图像转换为 RGB 图像，或者修改 VGG16 的第一层卷积层的输入通道数。
构建和编译模型:
- 实例化 VGG16TransferLearning 类，构建迁移学习模型。
- 使用 model.compile 方法配置模型的训练过程，包括损失函数、优化器和评估指标。
训练模型: 使用 model.fit 方法训练模型。

Kernel崩溃问题及解决方案

在训练过程中遇到Kernel崩溃的问题，通常与以下几个因素有关：

内存不足: 迁移学习模型通常比较大，训练时需要消耗大量的内存。如果计算机的内存不足，可能会导致Kernel崩溃。
GPU资源不足: 如果使用GPU加速训练，但GPU的显存不足，也可能导致Kernel崩溃。
TensorFlow配置问题: TensorFlow可能无法正确识别或使用GPU资源。

解决方案：

检查GPU是否可用： 确保TensorFlow能够检测到GPU。可以在Python中运行以下代码：
```
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
```
如果输出为0，则表示TensorFlow没有检测到GPU。需要检查TensorFlow和CUDA的安装是否正确，以及是否正确配置了环境变量。

限制GPU使用量： 如果TensorFlow能够检测到GPU，但仍然出现Kernel崩溃，可以尝试限制TensorFlow使用的GPU显存量。可以使用以下代码：

AI Content Detector

Writer推出的AI内容检测工具

下载

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)

这段代码启用了GPU的内存增长，允许TensorFlow根据需要动态分配显存，而不是一次性占用所有显存。

减小batch size： 减小batch size可以减少每次迭代所需的内存和显存，从而降低Kernel崩溃的风险。
检查输入数据： 确保输入数据的形状与模型的输入形状匹配。如果输入数据的形状不正确，可能会导致模型在训练过程中出现错误，从而导致Kernel崩溃。在上面的例子中，base_model的输入形状是(75, 75, 3)，因此需要确保x_train和x_test的形状也是(75, 75, 3)。如果MNIST数据集是灰度图像，需要将其转换为RGB图像，或者修改VGG16的第一层卷积层的输入通道数。
更新TensorFlow和CUDA驱动： 确保TensorFlow版本与CUDA和cuDNN版本兼容。使用最新版本的驱动程序通常可以解决一些兼容性问题。
简化模型： 如果以上方法都无法解决问题，可以尝试简化模型，例如减少全连接层的神经元数量，或者使用更小的预训练模型。

MNIST数据预处理

由于VGG16模型期望输入RGB图像，而MNIST数据集是灰度图像，因此需要对MNIST数据进行预处理，将其转换为RGB图像。可以使用以下代码：

import numpy as np
from tensorflow.keras.datasets import mnist

# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 将灰度图像转换为RGB图像
x_train = np.stack((x_train,) * 3, axis=-1)
x_test = np.stack((x_test,) * 3, axis=-1)

# 将图像大小调整为 (75, 75)
x_train = tf.image.resize(x_train, (75, 75))
x_test = tf.image.resize(x_test, (75, 75))

# 归一化像素值
x_train = x_train / 255.0
x_test = x_test / 255.0

print("x_train shape:", x_train.shape)
print("x_test shape:", x_test.shape)

这段代码首先加载MNIST数据集，然后使用 np.stack 将灰度图像复制三份，形成RGB图像。接着，使用 tf.image.resize 将图像大小调整为 (75, 75)，以匹配 VGG16 模型的输入形状。最后，将像素值归一化到 [0, 1] 范围内。

总结

本文介绍了如何使用VGG16模型进行MNIST手写数字的迁移学习，并重点讨论了在训练过程中可能遇到的Kernel崩溃问题。通过检查GPU配置、限制GPU使用量、减小batch size、检查输入数据、更新驱动程序和简化模型等方法，可以有效地解决Kernel崩溃问题，顺利完成模型的训练。此外，本文还介绍了MNIST数据的预处理方法，确保输入数据与VGG16模型的输入要求一致。希望本文能够帮助读者更好地理解和应用迁移学习技术。

Python 中实现单实例应用并等待前一个实例完成的完整方案

Python 中实现单实例应用并等待前一个实例结束的完整方案

如何高效统计字符串中出现的唯一排列子串数量

如何修改 Python 中的元组元素？

如何正确统计字符串中目标串的不重复排列子串数量