保姆级教程:用facenet-pytorch在Colab上训练自己的人脸识别模型(附数据集处理与避坑指南)

张开发
2026/4/20 19:50:34 15 分钟阅读

分享文章

保姆级教程:用facenet-pytorch在Colab上训练自己的人脸识别模型(附数据集处理与避坑指南)
零基础实战在Colab上构建高精度人脸识别模型的完整指南人脸识别技术正逐渐渗透到日常生活的各个角落从手机解锁到门禁系统这项技术正在重新定义身份验证的方式。对于开发者而言掌握构建自定义人脸识别模型的能力意味着可以打造更贴合特定场景的智能应用。本文将带领你从零开始在Google Colab的免费GPU环境中使用facenet-pytorch库完成一个完整的人脸识别项目。1. 环境准备与工具选择Google Colab作为云端Jupyter笔记本环境提供了免费的GPU计算资源特别适合深度学习模型的训练与实验。与本地环境相比Colab省去了复杂的驱动安装和环境配置步骤让开发者能够专注于模型本身。在开始之前确保你已经拥有Google账号并登录Colabhttps://colab.research.google.com/了解基本的Python语法和PyTorch框架准备好需要训练的人脸图片数据集至少包含3个不同人物的照片每人10-20张提示Colab的GPU资源有限制使用时长长时间训练建议保存检查点或考虑Colab Pro安装facenet-pytorch库的三种方式对比安装方式命令适用场景优缺点pip直接安装!pip install facenet-pytorch快速开始不需修改源码最简单但无法自定义修改GitHub克隆!git clone https://github.com/timesler/facenet-pytorch需要查看源码可查看但安装稍复杂手动安装下载源码后!pip install .需要修改库代码最灵活但步骤最多推荐初学者使用第一种方式执行以下命令即可!pip install facenet-pytorch import torch from facenet_pytorch import MTCNN, InceptionResnetV1 # 检查GPU是否可用 device torch.device(cuda:0 if torch.cuda.is_available() else cpu) print(f运行设备: {device})2. 数据集准备与预处理一个优质的数据集是模型成功的关键。对于人脸识别任务数据集需要满足以下要求每个类别人物至少10-20张照片图片包含清晰的正脸避免过度遮挡光照条件多样提高模型鲁棒性图像尺寸建议不低于160×160像素典型的数据集目录结构应如下/data/ /person1/ image1.jpg image2.jpg ... /person2/ image1.jpg image2.jpg ... ...使用MTCNN进行人脸检测和对齐的完整流程from torchvision import datasets from torch.utils.data import DataLoader import os # 设置数据集路径 data_dir ./data/your_dataset # 替换为你的数据集路径 batch_size 32 workers 0 if os.name nt else 4 # 初始化MTCNN mtcnn MTCNN( image_size160, margin0, min_face_size20, thresholds[0.6, 0.7, 0.7], factor0.709, post_processTrue, devicedevice ) # 准备数据集 dataset datasets.ImageFolder(data_dir, transformtransforms.Resize((512, 512))) dataset.samples [ (p, p.replace(data_dir, data_dir _cropped)) for p, _ in dataset.samples ] loader DataLoader( dataset, num_workersworkers, batch_sizebatch_size, collate_fncollate_fn ) # 执行人脸检测和对齐 for i, (x, y) in enumerate(loader): mtcnn(x, save_pathy) print(f\r处理批次 {i 1}/{len(loader)}, end) # 释放MTCNN内存 del mtcnn常见问题及解决方案检测不到人脸调整MTCNN的thresholds参数降低检测阈值对齐效果差增大margin值给检测框留出更多余量内存不足减小batch_size或使用Colab提供的更高内存运行时3. 模型训练与微调facenet-pytorch提供了在VGGFace2上预训练的InceptionResnetV1模型我们可以在此基础上进行微调。以下是完整的训练流程from torch.optim import Adam from torch.optim.lr_scheduler import MultiStepLR from torch.utils.tensorboard import SummaryWriter from torchvision import transforms import numpy as np # 初始化模型 resnet InceptionResnetV1( classifyTrue, pretrainedvggface2, num_classeslen(dataset.class_to_idx) ).to(device) # 定义优化器和学习率调度器 optimizer Adam(resnet.parameters(), lr0.001) scheduler MultiStepLR(optimizer, [5, 10]) # 数据预处理 trans transforms.Compose([ np.float32, transforms.ToTensor(), fixed_image_standardization ]) # 加载对齐后的数据集 dataset datasets.ImageFolder(data_dir _cropped, transformtrans) img_inds np.arange(len(dataset)) np.random.shuffle(img_inds) # 划分训练集和验证集80%-20% train_inds img_inds[:int(0.8 * len(img_inds))] val_inds img_inds[int(0.8 * len(img_inds)):] train_loader DataLoader( dataset, num_workersworkers, batch_sizebatch_size, samplerSubsetRandomSampler(train_inds) ) val_loader DataLoader( dataset, num_workersworkers, batch_sizebatch_size, samplerSubsetRandomSampler(val_inds) ) # 训练配置 loss_fn torch.nn.CrossEntropyLoss() metrics { fps: training.BatchTimer(), acc: training.accuracy } writer SummaryWriter() writer.iteration, writer.interval 0, 100 # 训练循环 for epoch in range(epochs): print(f\nEpoch {epoch 1}/{epochs}) print(- * 10) # 训练阶段 resnet.train() training.pass_epoch( resnet, loss_fn, train_loader, optimizer, scheduler, batch_metricsmetrics, show_runningTrue, devicedevice, writerwriter ) # 验证阶段 resnet.eval() training.pass_epoch( resnet, loss_fn, val_loader, batch_metricsmetrics, show_runningTrue, devicedevice, writerwriter ) writer.close()训练过程中的关键参数调整建议学习率初始0.001在第5和第10epoch时衰减批量大小根据GPU内存选择通常32-64之间训练轮数小型数据集8-15轮足够大型数据集可增加数据增强可添加随机翻转、旋转等增强模型泛化能力注意训练过程中可以使用TensorBoard监控指标在Colab中运行%load_ext tensorboard然后%tensorboard --logdir runs即可查看4. 模型评估与部署训练完成后我们需要评估模型性能并保存结果。以下是模型保存和预测的完整代码# 保存模型 model_path face_recognition_model.pth torch.save(resnet.state_dict(), model_path) # 加载模型进行预测 model InceptionResnetV1(classifyTrue, num_classeslen(dataset.class_to_idx)).to(device) model.load_state_dict(torch.load(model_path)) model.eval() # 定义预测函数 def predict_image(model, image_path, transform): test_image Image.open(image_path).convert(RGB) test_image_tensor transform(test_image) if torch.cuda.is_available(): test_image_tensor test_image_tensor.unsqueeze(0).cuda() else: test_image_tensor test_image_tensor.unsqueeze(0) with torch.no_grad(): outputs model(test_image_tensor) _, predicted torch.max(outputs, 1) return predicted.item() # 测试单张图片 test_transform transforms.Compose([ transforms.Resize(160), transforms.ToTensor(), fixed_image_standardization ]) class_names dataset.classes test_image_path ./data/test_images_cropped/person1/1.jpg # 替换为你的测试图片 predicted_class predict_image(model, test_image_path, test_transform) print(f预测结果: {class_names[predicted_class]})模型性能优化技巧测试集准确率计算添加验证集评估代码监控模型真实表现错误分析检查分类错误的样本找出模型弱点模型量化使用torch.quantization减小模型尺寸提升推理速度ONNX导出将模型转为ONNX格式便于跨平台部署实际项目中遇到的几个实用经验对于小数据集冻结底层特征提取层只训练顶层分类器效果更好人脸对齐步骤对最终准确率影响显著务必确保对齐质量不同光照条件下的测试能有效评估模型鲁棒性添加负样本非人脸或未知人物可以提高实际应用中的可靠性

更多文章