AI项目十一:Swin Transformer训练

若该文为原创文章,转载请注明原文出处。

续上一篇,训练自己的数据集,并测试。

一、安装标注软件labelme

# 安装labelme
pip install labelme
# 启动
labelme

AI项目十一:Swin Transformer训练_第1张图片

这里数据集准本,标注图片数据过程自己探索。

最后文件结构如下:

AI项目十一:Swin Transformer训练_第2张图片

二、修改配置文件

1、 修改configs\_base_\models\mask_rcnn_swin_fpn.py第54、73行num_classes为自己的类别数

我的类型是4个,所以填写为4

AI项目十一:Swin Transformer训练_第3张图片

2、运行 python modify.py 修改预训练模型

python .\modify.py --weights mask_rcnn_swin_tiny_patch4_window7_1x.pth --num_class 4 --output model_new.pt

3、修改configs\_base_\default_runtime.py,在最后增加一句加载预训练模型命令,用绝对路径

AI项目十一:Swin Transformer训练_第4张图片

4、修改configs\_base_\datasets\coco_instance.py第31-32行数据加载情况

AI项目十一:Swin Transformer训练_第5张图片

5、修改mmdet\datasets\coco.py第23行改为自己的标注,label顺序在coco_validate.ipynb中查看

CLASSES = ('arrow', 'car', 'dashed', 'line')

三、训练

python tools/train.py configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py

报错:AssertionError: Incompatible version of pycocotools is installed. Run pip uninstall pycocotools first. Then run pip install mmpycocotools to install open-mmlab forked pycocotools.

根据操作执行:

pip uninstall pycocotools

pip install mmpycocotools

 在次执行训练命令,可以看出已经在训练了,

AI项目十一:Swin Transformer训练_第6张图片

AI项目十一:Swin Transformer训练_第7张图片

训练结束后,模型在目录​​​work_dirs/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco下。

四、测试

根据上一篇测试命令,测试

python demo/video_demo.py 1.mp4 configs/swin/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco.py work_dirs/mask_rcnn_swin_tiny_patch4_window7_mstrain_480-800_adamw_1x_coco/latest.pth --out out.mp4

AI项目十一:Swin Transformer训练_第8张图片

如有侵权,或需要完整代码,请及时联系博主。

你可能感兴趣的:(AI计算机视觉,人工智能,transformer,深度学习)