使用VGG16进行MNIST手写数字识别的迁移学习-Python教程-PHP中文网

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

霞舞

发布： 2025-08-19 19:24:25

原创

909人浏览过

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

本文档旨在指导读者使用VGG16模型进行MNIST手写数字识别的迁移学习。我们将重点介绍如何构建模型、加载预训练权重、以及解决可能遇到的GPU配置问题。通过本文，读者可以掌握利用VGG16进行图像分类任务迁移学习的基本方法，并了解如何调试TensorFlow在GPU上的运行环境。

VGG16迁移学习实现MNIST手写数字识别

迁移学习是一种强大的技术，它允许我们利用在大规模数据集上预训练的模型，并将其应用于新的、通常较小的数据集。这可以显著减少训练时间和所需的计算资源，同时还能提高模型的性能。在本文中，我们将使用在ImageNet数据集上预训练的VGG16模型进行MNIST手写数字识别。

数据准备

首先，我们需要加载MNIST数据集并进行预处理。MNIST数据集包含 0 到 9 的手写数字的灰度图像。为了与 VGG16 兼容，我们需要将灰度图像转换为三通道图像，并将图像大小调整为 VGG16 期望的输入大小。

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

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

# 将灰度图像转换为三通道图像
x_train = np.stack([x_train]*3, axis=-1)
x_test = np.stack([x_test]*3, axis=-1)

# 调整图像大小为 (75, 75, 3)
img_height, img_width = 75, 75
x_train = tf.image.resize(x_train, (img_height, img_width)).numpy()
x_test = tf.image.resize(x_test, (img_height, img_width)).numpy()

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

登录后复制

构建VGG16迁移学习模型

接下来，我们将创建一个基于VGG16的迁移学习模型。我们将加载预训练的VGG16模型，移除其顶层（分类层），并添加我们自己的自定义分类层。

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):
    activation_list = []
    out = args[0]

    for layer in self.model.layers:
      out = layer(out)
      activation_list.append(out)
    if kwargs.get('training', False):
      return out
    else:
      prob = tf.nn.softmax(out)
      return out, prob

登录后复制

# 加载预训练的VGG16模型，不包含顶层
base_model = VGG16(weights="imagenet", include_top=False, input_shape=x_train[0].shape)

# 冻结VGG16模型的权重，防止在训练过程中被修改
base_model.trainable = False

# 创建迁移学习模型
model = VGG16TransferLearning(base_model, models)

登录后复制

编译和训练模型

现在，我们可以编译和训练我们的迁移学习模型。

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

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

登录后复制

GPU配置问题及解决

在训练过程中，可能会遇到Kernel Restarting的问题，这通常是由于GPU配置不正确导致的。即使你的电脑配备了GPU，TensorFlow也可能无法正确识别和使用它。

AI Sofiya

一款AI驱动的多功能工具

103

查看详情

以下是一些可能的解决方法：

检查TensorFlow是否检测到GPU：
```
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
```
登录后复制
如果输出为0，则表示TensorFlow没有检测到GPU。
安装正确的CUDA和cuDNN版本： TensorFlow需要特定版本的CUDA和cuDNN才能正常使用GPU。请参考TensorFlow官方文档，安装与你的TensorFlow版本兼容的CUDA和cuDNN版本。
设置环境变量： 确保CUDA和cuDNN的路径已添加到环境变量中。

使用tf.config.experimental.set_memory_growth： 在某些情况下，TensorFlow可能无法正确分配GPU内存。可以使用以下代码限制TensorFlow使用的GPU内存量：

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)

登录后复制

通过以上步骤，可以解决TensorFlow无法识别GPU的问题，从而避免Kernel Restarting错误。