基于预训练模型的图像分类实战:PyTorch与TensorFlow双框架实现
📌 前言 在实际的计算机视觉工程中,从零开始训练一个深度神经网络往往面临数据量不足与算力开销巨大的问题。迁移学习通过复用在大规模数据集(如ImageNet)上预训练的模型,将其特征提取能力迁移至特定任务,已成为工业界解决图像分类的标准范式。本文将基于预训练的ResNet/MobileNet模型,使用PyTorch与TensorFlow双框架,完整演示从数据预处理、模型微调到训练评估的迁移学习实战流程。
🧰 一、环境准备 确保Python版本为3.8+,需安装以下核心依赖包:
pip install torch torchvision tensorflow tensorflow-datasets matplotlib numpy
🖼️ 二、数据加载与预处理 迁移学习对预处理有严格要求,输入图像必须与预训练模型原有的归一化方式保持一致(如ImageNet的均值与标准差),且需将尺寸缩放至模型要求的224×224。
2.1 PyTorch 数据加载 使用torchvision构建数据流水线,应用Resize、标准化等变换。
import torch
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
# 定义ImageNet标准的预处理
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 以CIFAR10为例(实际业务替换为自定义ImageFolder数据集)
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
2.2 TensorFlow 数据加载 使用tensorflow_datasets加载自带数据,利用tf.data进行高效预处理与批处理。
import tensorflow as tf
import tensorflow_datasets as tfds
# 定义预处理函数
def preprocess(image, label):
image = tf.image.resize(image, (224, 224))
image = tf.cast(image, tf.float32) / 255.0
# 应用ImageNet归一化
image = (image - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
return image, label
# 加载tf_flowers数据集 (5分类)
dataset, info = tfds.load('tf_flowers', with_info=True, as_supervised=True)
train_ds = dataset['train'].take(3000).map(preprocess).shuffle(1000).batch(32)
test_ds = dataset['train'].skip(3000).map(preprocess).batch(32)
num_classes = 5
🏗️ 三、模型构建与微调 加载预训练模型,冻结卷积基网络权重以保留通用特征,仅替换并训练顶层分类器以适应新任务(微调阶段可解冻部分顶层卷积层)。
3.1 PyTorch 模型微调 基于ResNet18,冻结特征提取层,重写全连接层。
import torch.nn as nn
# 加载预训练ResNet18
model_pt = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# 冻结卷积基网络参数
for param in model_pt.parameters():
param.requires_grad = False
# 重写全连接层,CIFAR10为10分类
num_ftrs = model_pt.fc.in_features
model_pt.fc = nn.Linear(num_ftrs, 10)
3.2 TensorFlow 模型微调 基于MobileNetV2,设置trainable=False冻结基网络,顶层添加自定义分类层。
from tensorflow.keras import layers, Model
# 加载预训练MobileNetV2,不含顶层
base_model = tf.keras.applications.MobileNetV2(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
base_model.trainable = False # 冻结基网络
# 构建新顶层分类器
inputs = tf.keras.Input(shape=(224, 224, 3))
x = base_model(inputs, training=False)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dropout(0.2)(x)
outputs = layers.Dense(num_classes, activation='softmax')(x)
model_tf = Model(inputs, outputs)
🚀 四、模型训练 设置交叉熵损失与优化器。注意,微调阶段通常使用较小的学习率(如1e-3或更低),避免破坏预训练特征。
4.1 PyTorch 训练流程
criterion = nn.CrossEntropyLoss()
# 仅优化未冻结的全连接层参数
optimizer = torch.optim.Adam(model_pt.fc.parameters(), lr=1e-3)
model_pt.train()
for epoch in range(5):
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model_pt(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print(f'Epoch {epoch+1}/5, Loss: {loss.item():.4f}')
4.2 TensorFlow 训练流程
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
model_tf.compile(optimizer=optimizer, loss=loss_fn, metrics=['accuracy'])
model_tf.fit(train_ds, epochs=5)
🧪 五、模型评估 在测试集上验证模型迁移后的泛化能力。
5.1 PyTorch 评估
model_pt.eval()
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_loader:
outputs = model_pt(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'PyTorch 测试集准确率: {100 * correct / total:.2f}%')
5.2 TensorFlow 评估
test_loss, test_acc = model_tf.evaluate(test_ds, verbose=2)
print(f'TensorFlow 测试集准确率: {test_acc:.4f}')
📝 六、总结 迁移学习大幅降低了图像分类任务的落地门槛。核心流程可概括为:
- 数据对齐:预处理需严格匹配预训练模型的归一化参数与输入尺寸。
- 基网络冻结:初期冻结特征提取器,仅训练分类头,防止破坏通用特征。
- 小学习率微调:使用较小的学习率进行迭代,必要时可解冻顶层卷积层进一步微调。 掌握此范式,即可将开源预训练模型快速迁移至各类垂直领域的图像分类业务中。
更多推荐




所有评论(0)