用AI训练AI：制作一个简单的猫狗识别模型

原创

yichen

发布于 2024-03-23 15:04:45

3392

发布于 2024-03-23 15:04:45

文章被收录于专栏：陈冠男的游戏人生陈冠男的游戏人生

这是腾讯云加社区共创官的选题互换挑战赛，一搭眼看到了这个题目，因为之前写过自己制作数据集并训练，这无非是换个数据源进行训练而已，于是果断选择了这个题目

但是之前学的知识已经忘得差不多了，然后突发奇想，这种模型的训练应该是很入门的了，网上教程一大堆，既然网上有相关的知识，那大模型应该能够应付得了，于是乎决定用 AI 训练一个 AI

训练数据是个比较麻烦的事情，想要让 AI 能够识别猫和狗，首先你得给他足够多的图片，让他知道什么样的是猫，什么样的是狗，这个事情就得自己动手做了，大模型总不能咔咔给你生成上万张图片出来

起初考虑的是直接网上下图片，之前学习的时候是通过插件批量下载的，但考虑到训练猫狗识别模型应该很多人都做过了，那么互联网上应该有别人已经整理过的图片，于是乎找到一个包含三万多张图片的数据集：https://www.heywhale.com/mw/dataset/5d11bb1a38dc33002bd6f1f1

在下载到的数据集中，有一个 train.zip，里面存放了 25000 张猫和狗的图片，并且以名字的形式标记出来了

接下来就看大模型的发挥了（这里采用的是一个国内第三方搭建的 GPT-4 模型）

Prompt：我有一个文件夹名为 train，里面存放了 12500 张猫的照片和 12500 张狗的照片，可以通过文件名开头是 dog 还是 cat 来区分这张照片是狗还是猫，请帮我编写一段代码，采用 python 的深度学习框架训练一个识别猫狗的模型，并给出调用模型进行测试的示例

================ GPT 回答分割线 START ================

为了训练一个识别猫狗的模型，我们可以使用Python的深度学习库，如TensorFlow或PyTorch。这里，我将给出一个使用TensorFlow和Keras进行分类的简单示例。这个例子将会展示如何加载数据、构建一个简单的卷积神经网络（CNN）模型进行训练，以及如何测试模型。

首先，确保你已经安装了TensorFlow。如果没有安装，可以通过以下命令安装：

pip install tensorflow

接下来是训练模型的代码示例：

import os
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 设置数据路径
base_dir = 'path/to/your/train/'
train_dir = os.path.join(base_dir)

# 数据预处理
train_datagen = ImageDataGenerator(rescale=1./255, validation_split=0.2) # 将像素值缩放到0到1之间，并划分20%的数据作为验证集

train_generator = train_datagen.flow_from_directory(
        train_dir,
        target_size=(150, 150),  # 调整图片大小为150x150
        batch_size=20,
        class_mode='binary',  # 因为是二分类问题，所以使用binary模式
        subset='training')  # 设置为训练数据

validation_generator = train_datagen.flow_from_directory(
        train_dir,
        target_size=(150, 150),
        batch_size=20,
        class_mode='binary',
        subset='validation')  # 设置为验证数据

# 构建模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(150, 150, 3)),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(128, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')  # 二分类输出层
])

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

# 训练模型
history = model.fit(
      train_generator,
      steps_per_epoch=100,  # 每轮迭代训练的步数
      epochs=15,  # 训练轮数
      validation_data=validation_generator,
      validation_steps=50,  # 验证集的步数
      verbose=2)

# 保存模型
model.save('cat_vs_dog_model.h5')

为了使用训练好的模型进行预测，你可以使用以下代码:

from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import numpy as np

# 加载模型
model = load_model('cat_vs_dog_model.h5')

# 加载一张图片并调整大小
img = image.load_img('path/to/your/image.jpg', target_size=(150, 150))
img_tensor = image.img_to_array(img)  # 将图片转换为numpy数组
img_tensor = np.expand_dims(img_tensor, axis=0)  # 为批处理添加一个维度
img_tensor /= 255.  # 模型训练时，图片像素值被缩放到了0到1之间，这里也需要相同的预处理

# 预测图片
prediction = model.predict(img_tensor)
if prediction[0] > 0.5:
    print("It's a dog.")
else:
    print("It's a cat.")

这段代码提供了一个非常基础的示例，用于训练和测试一个猫狗识别模型。在实际应用中，你可能需要进行更复杂的数据增强、模型调整和优化等操作来提高模型的表现。

================ GPT 回答分割线 END ================

但实际 GPT 高估了我的电脑环境，因为后来根本没再做过深度学习这块的训练，我还需要额外安装几个库

pip install numpy
pip install pillow

在执行之前大致看了一下，感觉没有针对图片名进行打标签这类的操作，于是又问了它：我的文件夹中是图片，并不是数据集，你是否少了制作数据集的步骤

================ GPT 回答分割线 START ================

理解了，如果你的文件夹中直接存放着图片，并且你希望通过文件名来区分猫和狗的图片，那么我们可以稍微调整上述代码来应对这种情况。这里，我们将直接使用图片文件夹作为数据源，并通过文件名前缀来区分类别。不过，由于我们已经使用了 ImageDataGenerator 和 flow_from_directory 方法，这些方法实际上可以自动处理这种文件结构，只要我们正确地组织文件夹。

假设你的文件夹结构如下：

train/
    cat/
        cat001.jpg
        cat002.jpg
        ...
    dog/
        dog001.jpg
        dog002.jpg
        ...

这种情况下，ImageDataGenerator可以通过子文件夹名（cat和dog）自动标记图片的类别，所以不需要额外的步骤来从文件名中提取类别信息。

如果你的图片并没有按照这种方式组织，而是所有的图片都直接放在一个train文件夹中，并且是通过文件名来区分（例如cat001.jpg, dog001.jpg），那么在使用ImageDataGenerator之前，你需要首先将图片分别移动到以类别命名的子文件夹中。

你可以使用以下脚本来自动地创建这样的结构（如果还没有的话）：

import os
import shutil

# 源文件夹路径
source_dir = 'path/to/your/train/'

# 目标文件夹路径
target_dir = 'path/to/your/train_reorganized/'

# 创建目标文件夹，如果不存在
if not os.path.exists(target_dir):
    os.makedirs(target_dir)

# 分别为猫和狗创建文件夹
for category in ['cat', 'dog']:
    category_dir = os.path.join(target_dir, category)
    if not os.path.exists(category_dir):
        os.makedirs(category_dir)

# 遍历源文件夹，根据文件名前缀将文件移动到相应的子文件夹中
for filename in os.listdir(source_dir):
    if filename.startswith('cat'):
        shutil.move(os.path.join(source_dir, filename), os.path.join(target_dir, 'cat', filename))
    elif filename.startswith('dog'):
        shutil.move(os.path.join(source_dir, filename), os.path.join(target_dir, 'dog', filename))

这段代码首先检查目标文件夹中是否已经有了按类别组织的图片，如果没有，它会创建相应的子文件夹，并根据文件名前缀将图片归类。执行完这段脚本后，你就可以使用我之前提供的代码来加载数据、训练模型和进行预测了。这里要确保在之前代码中的base_dir变量设置为你的目标文件夹路径target_dir。

================ GPT 回答分割线 END ================

原来是代码里的某个函数，可以根据文件夹结构自己进行打标签这类的操作呀，那么只需要按照 GPT 给出的目录格式组织我的数据就可以啦，这里就不用给的脚本移动图片了，我建个文件夹复制粘贴的事情比用脚本方便多了哈哈哈

然后复制他的脚本，重命名为 train.py，运行脚本就正常进行训练啦

训练结束就得到了一个 200 多 MB 的模型

再用 GPT 的测试脚本随便从 test1 文件夹中照一张图片，运行得到正确的预测结果：

最后，在 GPT 的帮助下，我们再来理解一下训练脚本每一步都做了什么事（代码详细注释版）

import os
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator 
# 导入必要的库
# ImageDataGenerator 是 Keras 库中用于图像增强和预处理的工具，可以方便地从目录中加载数据、进行数据扩增等

# 设置数据路径
base_dir = 'train'
train_dir = os.path.join(base_dir)

# 数据预处理，创建一个 ImageDataGenerator 实例，用于数据预处理
# ImageDataGenerator 类是 TensorFlow 中 tf.keras.preprocessing.image 模块的一部分，用于实时地生成批量图像数据
# 这个类通过在训练过程中对图像进行实时的数据增强来提高模型的泛化能力。数据增强包括一系列的图像转换操作，比如旋转、缩放、平移、剪裁、翻转等，以及像素值的标准化等预处理步骤
train_datagen = ImageDataGenerator(rescale=1./255, validation_split=0.2) 
# rescale=1./255 表示将图像像素值从[0,255]缩放到[0,1]区间，这有助于模型学习，validation_split=0.2 表示从训练数据中划分 20% 作为验证集

# flow_from_directory 是 ImageDataGenerator 类的一个方法，它用于从文件夹路径中直接加载图像，并将它们作为深度学习模型的输入
# 这个方法非常适合处理文件夹中按类别组织的图像数据，它按照文件夹的结构自动为图像分配标签
train_generator = train_datagen.flow_from_directory(
        train_dir, # 从目录加载图像数据。为训练和验证数据设置了不同的参数，包括图像大小、批次大小、类别模式及数据子集类型
        target_size=(150, 150),  # 调整图片大小为150x150
        batch_size=20,
        class_mode='binary',  # 因为是二分类问题，所以使用binary模式
        subset='training')  # 设置为训练数据

validation_generator = train_datagen.flow_from_directory(
        train_dir,
        target_size=(150, 150),
        batch_size=20,
        class_mode='binary',
        subset='validation')  # 设置为验证数据

# 构建模型，使用Sequential模型来堆叠层，构建一个卷积神经网络（CNN）
# 该网络包含多个卷积层和池化层用于特征提取，一个展平层将二维图片数据转换为一维，以及两个密集层用于分类
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(150, 150, 3)),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(128, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')  # 二分类输出层
])

# 使用 adam 优化器、binary_crossentropy 损失函数（因为这是一个二分类问题）以及准确度评估指标来编译模型
model.compile(optimizer='adam',
              loss='binary_crossentropy',
              metrics=['accuracy'])

# 使用 fit 方法来训练模型，指定训练和验证数据生成器、步骤数、轮数和详细度参数
history = model.fit(
      train_generator,
      steps_per_epoch=100,  # 每轮迭代训练的步数
      epochs=15,  # 训练轮数
      validation_data=validation_generator,
      validation_steps=50,  # 验证集的步数
      verbose=2)

# 训练完成后，模型被保存为cat_vs_dog_model.h5
model.save('cat_vs_dog_model.h5')

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

gpt