图像分类是深度学习最基础也最广泛的应用方向之一,用Python训练图像分类模型需要覆盖数据准备、预处理、模型构建、训练、评估到最终预测的全流程。本文以TensorFlow框架为例,演示完整的实现步骤。

环境准备
首先需要安装必要的依赖库,核心库包括TensorFlow用于模型构建与训练,numpy用于数值处理,matplotlib用于结果可视化。执行以下命令安装依赖:
pip install tensorflow numpy matplotlib
数据集准备
这里使用TensorFlow内置的Fashion MNIST数据集作为示例,该数据集包含10个类别的服装灰度图像,共70000张,无需额外下载。如果是自定义数据集,需要将图像按类别存放在不同文件夹下,后续使用image_dataset_from_directory接口加载即可。
数据预处理
数据预处理主要包括归一化、划分训练集和验证集、构建数据管道三个步骤,代码如下:
import tensorflow as tf import numpy as np # 加载内置数据集 (train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data() # 归一化到0-1区间 train_images = train_images / 255.0 test_images = test_images / 255.0 # 调整图像维度,添加通道维度(灰度图为1通道) train_images = np.expand_dims(train_images, axis=-1) test_images = np.expand_dims(test_images, axis=-1) # 构建训练集和验证集的数据管道 train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).shuffle(10000).batch(32) test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(32)
构建卷积神经网络模型
卷积神经网络是图像分类任务的主流模型结构,这里构建一个包含卷积层、池化层和全连接层的基础模型:
from tensorflow.keras import layers, models
model = models.Sequential([
# 第一层卷积+池化
layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
layers.MaxPooling2D((2, 2)),
# 第二层卷积+池化
layers.Conv2D(64, (3, 3), activation='relu'),
layers.MaxPooling2D((2, 2)),
# 展平层
layers.Flatten(),
# 全连接层
layers.Dense(64, activation='relu'),
# 输出层,10个类别对应10个输出神经元,使用softmax激活
layers.Dense(10, activation='softmax')
])
# 查看模型结构
model.summary()
模型编译与训练
编译模型时需要指定优化器、损失函数和评估指标,这里使用Adam优化器,稀疏分类交叉熵损失,准确率作为评估指标:
# 编译模型
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy']
)
# 训练模型,训练10个epoch
history = model.fit(
train_dataset,
epochs=10,
validation_data=test_dataset
)
模型评估与预测
训练完成后,可以在测试集上评估模型的性能,同时使用训练好的模型对新图像进行预测:
import matplotlib.pyplot as plt
# 评估模型
test_loss, test_acc = model.evaluate(test_dataset)
print(f"测试集准确率: {test_acc:.4f}")
# 取测试集前5张图像进行预测
sample_images = test_images[:5]
sample_labels = test_labels[:5]
predictions = model.predict(sample_images)
# 可视化预测结果
class_names = ['T恤', '裤子', '套头衫', '连衣裙', '外套', '凉鞋', '衬衫', '运动鞋', '包', '短靴']
plt.figure(figsize=(10, 5))
for i in range(5):
plt.subplot(1, 5, i+1)
plt.imshow(sample_images[i].squeeze(), cmap='gray')
plt.title(f"真实: {class_names[sample_labels[i]]}n预测: {class_names[np.argmax(predictions[i])]}")
plt.axis('off')
plt.show()
自定义数据集适配说明
如果使用自己的图像数据集,只需要替换数据加载部分的代码即可,例如数据集存放在data/train和data/val目录下,每个类别对应一个子文件夹,可以使用以下方式加载:
train_dataset = tf.keras.preprocessing.image_dataset_from_directory(
'data/train',
image_size=(224, 224), # 调整图像尺寸
batch_size=32
)
val_dataset = tf.keras.preprocessing.image_dataset_from_directory(
'data/val',
image_size=(224, 224),
batch_size=32
)
# 获取类别名称
class_names = train_dataset.class_names
后续模型构建时需要注意调整输入形状匹配图像尺寸,输出层神经元数量匹配自定义数据集的类别数即可。