从数据集准备到模型训练：一步步教你用P2PNet跑通SHHA人群计数数据集

从数据集准备到模型训练一步步教你用P2PNet跑通SHHA人群计数数据集人群计数技术在智慧城市、公共安全等领域具有广泛应用价值。P2PNet作为近年来提出的创新性人群计数模型以其端到端的点预测能力和轻量化结构受到研究者关注。本文将带您从零开始完整实现P2PNet在SHHA数据集上的训练流程涵盖数据准备、环境配置、训练调参等关键环节特别针对实际项目中容易遇到的路径配置、版本兼容等问题提供解决方案。1. 环境准备与数据获取1.1 基础环境配置P2PNet基于PyTorch框架实现推荐使用Python 3.8和CUDA 11.3以上环境。以下是核心依赖的安装命令pip install torch2.1.2 torchvision0.16.2 pip install opencv-python pandas matplotlib tensorboardX常见环境问题排查CUDA版本不匹配通过nvidia-smi和nvcc --version检查驱动与运行时版本Pillow版本问题新版Pillow中ANTIALIAS已被移除需替换为Resampling.LANCZOSTorchvision API变更若遇到_new_empty_tensor导入错误可注释相关代码或降级torchvision版本1.2 SHHA数据集准备SHHAShanghaiTech Part A是人群计数领域的基准数据集包含482张高密度人群图像。建议按以下结构组织数据SHHA/ ├── images/ │ ├── IMG_1.jpg │ └── ... ├── txt/ │ ├── GT_IMG_1.txt │ └── ... └── splits/ ├── train.list └── test.list每个GT文件包含对应图像的标注点坐标格式示例104.5 203.2 78.1 156.7 ...提示原始数据集可能使用.mat格式标注需提前转换为txt格式。可使用以下Python代码片段转换import scipy.io mat scipy.io.loadmat(GT_IMG_1.mat) points mat[image_info][0][0][0][0][0] # SHHA特定结构 np.savetxt(GT_IMG_1.txt, points, fmt%.1f)2. 数据预处理实战2.1 自动生成数据列表文件P2PNet训练需要提供包含图像-标注对路径的列表文件。以下脚本可自动生成训练/测试集列表import os def generate_data_list(dataset_path, output_file): image_files [f for f in os.listdir(f{dataset_path}/images) if f.endswith(.jpg)] with open(output_file, w) as f: for img_file in image_files: img_path os.path.join(dataset_path, images, img_file) txt_file fGT_{os.path.splitext(img_file)[0]}.txt txt_path os.path.join(dataset_path, txt, txt_file) if os.path.exists(txt_path): f.write(f{img_path} {txt_path}\n) else: print(fWarning: Missing annotation for {img_file}) # 示例用法 generate_data_list(/data/SHHA/train, train.list) generate_data_list(/data/SHHA/test, test.list)关键参数说明dataset_path包含images和txt子目录的数据集根目录output_file生成的列表文件路径如train.list2.2 数据增强策略P2PNet原始论文采用以下增强组合随机水平翻转p0.5颜色抖动亮度0.2对比度0.2饱和度0.2随机裁剪512×512可通过修改datasets.py中的__getitem__方法调整增强策略。推荐保留原始增强方案以获得最佳性能。3. 模型训练全流程3.1 训练参数解析P2PNet的核心训练命令如下python train.py \ --data_root /path/to/SHHA \ --dataset_file SHHA \ --epochs 3500 \ --lr_drop 3500 \ --batch_size 8 \ --lr 0.0001 \ --lr_backbone 0.00001 \ --eval_freq 1 \ --output_dir ./logs \ --checkpoints_dir ./weights关键参数说明参数推荐值作用lr1e-4主学习率lr_backbone1e-5Backbone学习率lr_drop3500学习率衰减epocheval_freq1每N个epoch验证一次batch_size8根据GPU显存调整3.2 训练监控与调优通过TensorBoard可实时监控训练过程tensorboard --logdir./logs --port6006重点观察指标train_loss应平稳下降最终收敛到0.2-0.3val_mae验证集平均绝对误差SHHA上通常能达到60-70lr学习率变化曲线常见训练问题处理Loss震荡减小batch_size或降低学习率过拟合增加数据增强或提前停止训练显存不足减小batch_size或使用梯度累积4. 模型推理与部署4.1 测试集评估使用训练好的模型进行评估python run_test.py \ --weight_path ./weights/best.pth \ --output_dir ./results \ --dataset_path /path/to/SHHA/test输出包括预测密度图_density.jpg点预测结果_points.txt可视化标注_vis.jpg4.2 自定义数据推理对新的图像进行预测from models import build_model from PIL import Image model build_model(args) checkpoint torch.load(weights/best.pth, map_locationcpu) model.load_state_dict(checkpoint[model]) img Image.open(custom_image.jpg).convert(RGB) points model.predict(img) # 获取预测点坐标4.3 性能优化技巧TorchScript导出将模型转换为TorchScript提升推理速度traced_model torch.jit.trace(model, example_input) traced_model.save(p2pnet.pt)ONNX转换支持跨平台部署torch.onnx.export(model, dummy_input, p2pnet.onnx)TensorRT加速对CUDA核心进行优化可获得2-3倍速度提升实际部署中发现输入分辨率对推理速度影响显著。将图像缩放至800×600左右可在精度和速度间取得较好平衡。