项目简介
本项目是基于Python和PyTorch构建的图像分类深度学习系统,聚焦于运用Swin Transformer V2模型开展图像分类任务。它支持分布式训练、模型验证以及多图特征融合,适用于大规模图像数据的处理与分析。
项目的主要特性和功能
主要特性
- 借助Horovod库实现多GPU分布式训练,大幅提升训练速度。
- 支持Swin Transformer V2等先进模型,且允许自定义网络架构。
- 支持加载预训练模型权重,方便进行模型微调。
- 采用CosineLRScheduler进行动态学习率调整,优化训练过程。
- 提供模型验证脚本,可计算准确率、PR曲线等评价指标。
- 利用TensorBoardX记录训练日志,便于实时监控和后续分析。
功能
- 提供训练脚本,支持分布式训练、预训练模型加载、学习率调度等功能。
- 提供模型验证脚本,用于评估模型性能,输出准确率、PR曲线等指标。
- 使用TensorBoardX记录训练过程中的损失、准确率等信息。
- 支持四图特征融合训练,提升模型在复杂场景下的分类性能。
安装使用步骤
假设用户已经下载了本项目的源码文件。
1. 安装依赖
首先,确保已安装必要的Python库:
bash
pip install torch horovod tensorboardX
2. 运行训练脚本
使用命令行参数运行训练脚本,配置训练参数:
bash
horovodrun -np 8 python hago/train_hago_timm.py \
--list-file "/data1/liuyidi/scene_cls/V4.1.2/list/train_list/train_list_remove4_new_resample.txt" \
--root-dir "/data1/liuyidi/scene_cls/6b_2/" \
--batch-size 8 \
--img-size 256 \
--base-lr 2e-5 \
--net "timm.swinv2_base_window12to16_192to256_22kft1k" \
--sam 0 \
--mixup 0.8 \
--ckpt-log-dir '/data1/liuyidi/scene_cls/V4.1.2/log_dir/V4.1.2_test37_38_remove_act=8' \
--pretrained-ckpt '/data1/liuyidi/model/swin_V2/swinv2_base_patch4_window12to16_192to256_22kto1k_ft.pth' \
--final-drop 0.2 \
--rand-corner 1 \
--label_smoothing 0.1 \
--weight-decay 1e-8 \
--ckpt-save-interval 4000 \
--optimizer 'adamw' \
--accumulation_steps 16 \
--cos_step 1024000 \
--base_batchsize 1024 \
--warmup_step 64000
3. 运行验证脚本
使用验证脚本评估模型性能:
bash
DIRNAME=/data1/liuyidi/scene_cls/V4.1.2/log_dir/V4.1.2_test37_38_remove_swin_base_22kft1k
Nckpts=512000
Ninterval=8000
for ((i=512000; i<=$Nckpts; i+=$Ninterval)); do
part_n=`printf "%06d" ${i}`
horovodrun -np 7 python hago/predict_hago_timm.py \
--list-file "/data1/liuyidi/scene_cls/V4.1.2/list/train_list/val_merge.txt" \
--root-dir "/data1/liuyidi/scene_cls/6b_2/" \
--load-npy 0 \
--num-classes 43 \
--net "timm.swinv2_base_window12to16_192to256_22kft1k" \
--img-size 256 \
--batch-size 64 \
--ckpt "$DIRNAME/ckpt/checkpoint-iter-${part_n}.pyth" \
--out-per-n 10000 \
--out "$DIRNAME/val/val-$part_n"
python hago/val_hago.py \
"/data1/liuyidi/scene_cls/V4.1.2/list/train_list/val_merge.txt" \
"$DIRNAME/val/val-$part_n" \
"${DIRNAME}/log/val/" \
${i} &
done
4. 查看日志
使用TensorBoardX查看训练日志:
bash
tensorboard --logdir=/data1/liuyidi/scene_cls/V4.1.2/log_dir/V4.1.2_test37_38_remove_act=8
通过以上步骤,您可以顺利进行基于Swin Transformer V2的图像分类任务的训练与评估。
下载地址
点击下载 【提取码: 4003】【解压密码: www.makuang.net】