【Tensorflow 2】使用迁移学习（MobileNet）完成图像分类任务

最新推荐文章于 2026-04-23 07:43:23 发布

原创

最新推荐文章于 2026-04-23 07:43:23 发布 · 5.5k 阅读

标签

#tensorflow #深度学习 #卷积神经网络

本文介绍了如何使用Tensorflow的Keras API进行迁移学习，特别是利用预训练的MobileNetV2模型对图像进行分类。首先，通过`image_dataset_from_directory`将图片数据转换为dataset，然后进行数据增强以提高模型泛化能力。接着，构建基于MobileNetV2的基础模型，并对其进行微调。在训练后，模型在验证集上的准确率达到了97.67%。最后，展示了模型在测试集上的优秀效果，并对额外的测试数据集`test2`进行了预测。

前言：图像分类是CV领域相对比较成熟的一种技术，一般从基础开始学就是Tensorflow——>CNN——>手撸ResNet等算法——>迁移学习。
Tensorflow.keras.applications中有很多迁移学习的算法，只需要加载后下载参数，然后fine_tune稍微训练最后几层，就可以获得非常不错的效果。
本文主要是通过一系列代码指导大家如何完成迁移学习的使用。

一、导入数据，制作dataset

因为是图片，所以我们首先就是需要把图片转换成Tensorflow能理解的向量形式。另外因为图片的数量多了以后，处理效率会变慢，所以一般都要做成数据集（dataset），然后设置batch_size(32或64)。
如上，一般有2种做法：
第一种是通过tf.io.read_file(文件路径)，然后tf.image.decode_jpeg解析成向量，最后通过tf.data.Datasets.from_tensor_slices()转化成dataset，这种方式比较麻烦，但是稳定；
第二种是直接使用tf.keras.preprocessing.image_dataset_from_directory加载，这种方式的好处是：如果训练集或验证集中的图片是分类存放在不同的文件夹内，且文件夹名称就是类别（或可标识），那么通过这种方式就可以直接转成dataset，非常方便。

import tensorflow as tf
from tensorflow.keras.preprocessing import image_dataset_from_directory
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
import warnings
warnings.filterwarnings('ignore')

#数据所在文件夹
base_dir = './dataset'

train_dir = os.path.join(base_dir, 'train')
valid_dir = os.path.join(base_dir, 'valid')
test_dir = os.path.join(base_dir, 'test')
test2_dir = os.path.join(base_dir, 'test2')

BATCH_SIZE = 32
IMG_SIZE = (160, 160)

#image_dataset_from_directory根据目录结构自动分类，工业上比较可以迭代
train_dataset = image_dataset_from_directory(train_dir,
                                            batch_size=BATCH_SIZE,
                                            image_size=IMG_SIZE)

valid_dataset = image_dataset_from_directory(valid_dir,
                                            batch_size=BATCH_SIZE,
                                            image_size=IMG_SIZE)

test_dataset = image_dataset_from_directory(test_dir,
                                            batch_size=BATCH_SIZE,
                                            shuffle=False, #默认是True，最后核对会比较麻烦
                                            image_size=IMG_SIZE)

#【重要】获取分类名称，这个最后会用到
class_names = train_dataset.class_names
print(class_names)

#开辟一个固定内存，动态和batch_size一致
train_dataset = train_dataset.prefetch(buffer_size=BATCH_SIZE)
valid_dataset = valid_dataset.prefetch(buffer_size=BATCH_SIZE)
test_dataset = test_dataset.prefetch(buffer_size=BATCH_SIZE)