06.基于VGG-16算法Pytorch实现人脸识别
- 🍨 本文为🔗365天深度学习训练营中的学习记录博客
- 🍖 原作者:K同学啊
🏡 我的环境:使用01中创建的虚拟环境mnist
- 虚拟环境:mnist
Python 3.10.19
Name: torch, Version: 2.10.0+cu130
Name: torchvision,Version: 0.25.0+cu130
- 编译器:Positron
- 深度学习环境:Pytorch
这次复现的项目,整体目标还是图像分类, 这一次走的是迁移学习(Transfer Learning)的路线。 直接调用已经在大规模数据集上训练好的 VGG16 预训练模型,然后把它最后一层改成适合我们当前任务的输出类别数,再用自己的数据去训练。
完整流程
1)准备环境
导入 PyTorch、torchvision、图像处理等库,自动检测 GPU/CPU。
2)读取数据
用 ImageFolder 读取整个 48-data 文件夹中的图片数据,同时自动建立类别与数字标签的映射。
3)图像预处理
统一缩放到 224×224,转成张量,并按 ImageNet 的均值和标准差做标准化,这一步是为了匹配 VGG16 预训练模型的输入习惯。
4)划分训练集和测试集
不是提前手动分文件夹,而是用 random_split 按 8:2 随机划分。
5)构建 DataLoader
按 batch 读取图像数据,方便后续训练和测试。
6)加载预训练 VGG16
直接调用已经训练好的 VGG16 模型,不再从头搭 CNN。
7)冻结前面参数,修改最后分类层
保留 VGG16 已经学到的通用视觉特征,只让最后一层适应当前任务的类别数。
8)定义训练与测试函数
训练函数负责前向传播、计算损失、反向传播和更新参数;测试函数只做前向推理和指标统计,不更新参数。
9)设置优化器和学习率调度器
使用 SGD 优化器和 LambdaLR 学习率衰减策略,让训练过程更平稳。
10)开始训练
每个 epoch 先训练,再测试,同时记录指标;如果当前测试准确率超过历史最好结果,就深拷贝保存为 best_model。
11)保存最佳模型
训练结束后,把测试表现最好的模型参数保存下来。
12)绘图观察结果
画出训练集和测试集的准确率、损失曲线,判断模型是否收敛、是否过拟合。
13)单张图片预测
读取一张新图片,执行和训练时一样的预处理,然后输入模型得到分类结果。
一、 前期准备
1. 设置GPU
如果设备上支持GPU就使用GPU,否则使用CPU
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision
from torchvision import transforms, datasets
import os,PIL,pathlib,warnings
warnings.filterwarnings("ignore") #忽略警告信息
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device
2. 导入数据
数据来源需要自行下载
import os,PIL,random,pathlib
data_dir = './48-data/'
data_dir = pathlib.Path(data_dir)
data_paths = list(data_dir.glob('*'))
classeNames = [str(path).split("\\")[1] for path in data_paths]
classeNames
# 关于transforms.Compose的更多介绍可以参考:https://blog.csdn.net/qq_38251616/article/details/124878863
train_transforms = transforms.Compose([
transforms.Resize([224, 224]), # 将输入图片resize成统一尺寸
# transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.ToTensor(), # 将PIL Image或numpy.ndarray转换为tensor,并归一化到[0,1]之间
transforms.Normalize( # 标准化处理-->转换为标准正太分布(高斯分布),使模型更容易收敛
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]) # 其中 mean=[0.485,0.456,0.406]与std=[0.229,0.224,0.225] 从数据集中随机抽样计算得到的。
])
total_data = datasets.ImageFolder("./48-data/",transform=train_transforms)
total_data
total_data.class_to_idx
这一段前面项目中的图像预处理思路是一致的,只不过这次只有一套 train_transforms,因为整个数据先统一读取,再随机切分训练集和测试集。
这几步分别在干什么
(1)Resize([224, 224])
把所有图片统一缩放成 224×224。
这一步很重要,因为 VGG16 的输入通常就是这个尺寸,或者至少需要固定尺寸输入。
神经网络不喜欢每张图片大小都不一样,所以必须先统一。
(2)ToTensor()
把图片从 PIL 图像转成 PyTorch 的张量格式。
同时它还会把像素值从原来的 0~255 映射到 0~1。
(3)Normalize(mean, std)
做标准化,让数据分布更稳定,更有利于模型收敛。
这里的均值和标准差:
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
3. 划分数据集
train_size = int(0.8 * len(total_data))
test_size = len(total_data) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(total_data, [train_size, test_size])
train_dataset, test_dataset
batch_size = 32
train_dl = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
test_dl = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
for X, y in test_dl:
print("Shape of X [N, C, H, W]: ", X.shape)
print("Shape of y: ", y.shape, y.dtype)
break
同前面项目一样,不在赘述
二、调用官方的VGG-16模型
from torchvision.models import vgg16
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using {} device".format(device))
# 加载预训练模型,并且对模型进行微调
model = vgg16(pretrained = True).to(device) # 加载预训练的vgg16模型
for param in model.parameters():
param.requires_grad = False # 冻结模型的参数,这样子在训练的时候只训练最后一层的参数
# 修改classifier模块的第6层(即:(6): Linear(in_features=4096, out_features=2, bias=True))
# 注意查看我们下方打印出来的模型
model.classifier._modules['6'] = nn.Linear(4096,len(classeNames)) # 修改vgg16模型中最后一层全连接层,输出目标类别个数
model.to(device)
model
什么叫 pretrained=True
意思是:
加载一个已经在 ImageNet 大规模数据集 上训练好的 VGG16 模型参数,而不是随机初始化。
这相当于这个模型已经“见过很多图片”,学会了很多基础视觉特征,例如:
- 边缘
- 纹理
- 颜色组合
- 局部形状
- 更高层次的图像模式
所以当我们自己的数据量不是特别大时,直接用预训练模型通常比从头训练效果更好,也更快收敛。
for param in model.parameters():
param.requires_grad = False
这一段表示:
把 VGG16 原本所有参数先冻结住,训练时不更新它们。
为什么要冻结
因为前面的卷积层已经在 ImageNet 上学到了很强的通用图像特征。
如果您当前的数据集比较小,直接把全部参数都打开训练,容易:
- 训练慢
- 过拟合
- 把原本学到的通用特征破坏掉
所以这份代码的策略是:
- 前面的大网络参数保留
- 只改最后一层分类器来适应新任务
这是一种非常经典的迁移学习写法。
model.classifier._modules['6'] = nn.Linear(4096,len(classeNames))
model.to(device)
model
VGG16 原本最后一层是为了 ImageNet 1000 分类设计的,输出维度是 1000。
但我们当前任务不是 1000 分类,而是自己的类别数,所以必须把最后一层改掉。
三、 训练模型
1. 编写训练函数
# 训练循环
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset) # 训练集的大小
num_batches = len(dataloader) # 批次数目, (size/batch_size,向上取整)
train_loss, train_acc = 0, 0 # 初始化训练损失和正确率
for X, y in dataloader: # 获取图片及其标签
X, y = X.to(device), y.to(device)
# 计算预测误差
pred = model(X) # 网络输出
loss = loss_fn(pred, y) # 计算网络输出和真实值之间的差距,targets为真实值,计算二者差值即为损失
# 反向传播
optimizer.zero_grad() # grad属性归零
loss.backward() # 反向传播
optimizer.step() # 每一步自动更新
# 记录acc与loss
train_acc += (pred.argmax(1) == y).type(torch.float).sum().item()
train_loss += loss.item()
train_acc /= size
train_loss /= num_batches
return train_acc, train_loss
同前面项目基本一致
2. 编写测试函数
def test (dataloader, model, loss_fn):
size = len(dataloader.dataset) # 测试集的大小
num_batches = len(dataloader) # 批次数目, (size/batch_size,向上取整)
test_loss, test_acc = 0, 0
# 当不进行训练时,停止梯度更新,节省计算内存消耗
with torch.no_grad():
for imgs, target in dataloader:
imgs, target = imgs.to(device), target.to(device)
# 计算loss
target_pred = model(imgs)
loss = loss_fn(target_pred, target)
test_loss += loss.item()
test_acc += (target_pred.argmax(1) == target).type(torch.float).sum().item()
test_acc /= size
test_loss /= num_batches
return test_acc, test_loss
同前面项目基本一致
3. 设置动态学习率
# def adjust_learning_rate(optimizer, epoch, start_lr):
# # 每 2 个epoch衰减到原来的 0.98
# lr = start_lr * (0.92 ** (epoch // 2))
# for param_group in optimizer.param_groups:
# param_group['lr'] = lr
learn_rate = 1e-4 # 初始学习率
# optimizer = torch.optim.SGD(model.parameters(), lr=learn_rate)
# 调用官方动态学习率接口时使用
lambda1 = lambda epoch: 0.92 ** (epoch // 4)
optimizer = torch.optim.SGD(model.parameters(), lr=learn_rate)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1) #选定调整方法
不再手写 adjust_learning_rate() 函数,而是改成用官方的学习率调度器
4. 正式训练
import copy
loss_fn = nn.CrossEntropyLoss() # 创建损失函数
epochs = 40
train_loss = []
train_acc = []
test_loss = []
test_acc = []
best_acc = 0 # 设置一个最佳准确率,作为最佳模型的判别指标
for epoch in range(epochs):
# 更新学习率(使用自定义学习率时使用)
# adjust_learning_rate(optimizer, epoch, learn_rate)
model.train()
epoch_train_acc, epoch_train_loss = train(train_dl, model, loss_fn, optimizer)
scheduler.step() # 更新学习率(调用官方动态学习率接口时使用)
model.eval()
epoch_test_acc, epoch_test_loss = test(test_dl, model, loss_fn)
# 保存最佳模型到 best_model
if epoch_test_acc > best_acc:
best_acc = epoch_test_acc
best_model = copy.deepcopy(model)
train_acc.append(epoch_train_acc)
train_loss.append(epoch_train_loss)
test_acc.append(epoch_test_acc)
test_loss.append(epoch_test_loss)
# 获取当前的学习率
lr = optimizer.state_dict()['param_groups'][0]['lr']
template = ('Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%, Test_loss:{:.3f}, Lr:{:.2E}')
print(template.format(epoch+1, epoch_train_acc*100, epoch_train_loss,
epoch_test_acc*100, epoch_test_loss, lr))
# 保存最佳模型到文件中
PATH = './best_model.pth' # 保存的参数文件名
torch.save(best_model.state_dict(), PATH)
print('Done')
import copy
loss_fn = nn.CrossEntropyLoss()
epochs = 40
train_loss = []
train_acc = []
test_loss = []
test_acc = []
best_acc = 0
这一段是在训练正式开始前做准备:
loss_fn:损失函数epochs=40:总训练轮数- 四个列表:记录每轮训练和测试的 loss、acc
best_acc=0:用于记录当前出现过的最高测试准确率
forepochinrange(epochs):
model.train()
epoch_train_acc, epoch_train_loss=train(train_dl, model, loss_fn, optimizer)
scheduler.step()
model.eval()
epoch_test_acc, epoch_test_loss=test(test_dl, model, loss_fn)
ifepoch_test_acc>best_acc:
best_acc=epoch_test_acc
best_model=copy.deepcopy(model)
train_acc.append(epoch_train_acc)
train_loss.append(epoch_train_loss)
test_acc.append(epoch_test_acc)
test_loss.append(epoch_test_loss)
lr=optimizer.state_dict()['param_groups'][0]['lr']
template= ('Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%, Test_loss:{:.3f}, Lr:{:.2E}')
print(template.format(epoch+1, epoch_train_acc*100, epoch_train_loss,
epoch_test_acc*100, epoch_test_loss, lr))
这一段可以理解成:每个 epoch 都做一次完整的训练+测试+记录+保存最佳模型。
关键步骤解释
model.train()
把模型切换到训练模式。
有些层在训练和测试时行为不同,比如 Dropout、BatchNorm。
虽然 VGG16 这里主要训练最后一层,但规范写法仍然要加。
train(...)
执行一个 epoch 的训练。
scheduler.step()
更新学习率。
这一步表示每轮训练后,让学习率按预定规则变化。
model.eval()
切换到评估模式。
让模型在测试时关闭训练态的特殊行为。
test(...)
在测试集上评估当前模型性能。
四、 结果可视化
1. Loss与Accuracy图
import matplotlib.pyplot as plt
#隐藏警告
import warnings
warnings.filterwarnings("ignore") #忽略警告信息
plt.rcParams['font.sans-serif'] = ['SimHei'] # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False # 用来正常显示负号
plt.rcParams['figure.dpi'] = 100 #分辨率
from datetime import datetime
current_time = datetime.now() # 获取当前时间
epochs_range = range(epochs)
plt.figure(figsize=(12, 3))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_acc, label='Training Accuracy')
plt.plot(epochs_range, test_acc, label='Test Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')
plt.xlabel(current_time) # 打卡请带上时间戳,否则代码截图无效
plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss, label='Training Loss')
plt.plot(epochs_range, test_loss, label='Test Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()
和之前项目一样

2. 指定图片进行预测
from PIL import Image
classes = list(total_data.class_to_idx)
def predict_one_image(image_path, model, transform, classes):
test_img = Image.open(image_path).convert('RGB')
plt.imshow(test_img) # 展示预测的图片
test_img = transform(test_img)
img = test_img.to(device).unsqueeze(0)
model.eval()
output = model(img)
_,pred = torch.max(output,1)
pred_class = classes[pred]
print(f'预测结果是:{pred_class}')
# 预测训练集中的某张照片
predict_one_image(image_path='./48-data/Angelina Jolie/010_f99d79e3.jpg',
model=model,
transform=train_transforms,
classes=classes)

3. 模型评估
best_model.eval()
epoch_test_acc, epoch_test_loss = test(test_dl, best_model, loss_fn)
epoch_test_acc, epoch_test_loss
# 查看是否与我们记录的最高准确率一致
epoch_test_acc
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐

所有评论(0)