Pytorch作为一个较新的开源框架,十分简洁好用,完全不亚于Tensorflow等成熟框架。
最近在学习Faster R-CNN, 发现Pytorch版本的资料不多,所以在这里记录与分享下自己安装配置Pytorch版本的faster cnn的过程。过程是搬运+翻译的说明文档,增加了一些细节性的东西和报错解决。
这里我用的是github上ruotianluo开源的代码:pytorch-faster-rcnn
一、安装
- 下载代码
git clone https://github.com/ruotianluo/pytorch-faster-rcnin.git
- 编译模块(nms, roi_pooling, roi_align)
这里可能会报错,后面讲解决方案
cd pytorch-faster-rcnn/lib
bash make.sh
cd ../
- 安装python COCO API (获取COCO数据集用)
cd data
git clone https://github.com/pdollar/coco.git cd coco/PythonAPI
make
cd ../../..
编译的过程中可能会报错,这时候对照下这个表,去 make.sh 里删掉的不支持的arch
GPU model | Architecture |
---|---|
TitanX (Maxwell/Pascal) | sm_52 |
GTX 960M | sm_50 |
GTX 1080 (Ti) | sm_61 |
Grid K520 (AWS g2.2xlarge) | sm_30 |
Tesla K80 (AWS p2.xlarge) | sm_37 |
我在编译的时候也疯狂报错,最后是进入make.sh里把-gencode arch=compute_70,code=sm_70 这一行删了就Ok了
二、准备数据集
这个请根据rbg大神的说明来操作 数据集下载与配置
这里直接搬运+翻译下
- 下载训练、验证、测试数据集和VOCdevkit
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCdevkit_08-Jun-2007.tar
- 将3个压缩包解压到该文件夹下
tar xvf VOCtrainval_06-Nov-2007.tar
tar xvf VOCtest_06-Nov-2007.tar
tar xvf VOCdevkit_08-Jun-2007.tar
- 解压完成后应呈如下结构
# 这里的$VOCdekit代指大家电脑上各自的真实目录
$VOCdevkit/ # 开发组件
$VOCdevkit/VOCcode/ # 工具类代码
$VOCdevkit/VOC2007 # 图片数据集等
# ... 其他的一些文件夹
- 为PASCAL VOC数据集创建软链接
# $FRCN_ROOT指是faster rcnn的根目录,这里是指代,大家使用时改成自己的地址
cd $FRCN_ROOT/data
ln -s $VOCdevkit VOCdevkit2007
来自官方的推荐:使用软链接非常方便,可以让同个数据集在不同项目之间共享
也就是说,你的数据集可以放在任意目录,然后在data文件夹下创建一个软链接(类似快捷方式)指向这个数据集
其他数据集的下载配置也类似,大家可以参考rbg大神的说明文档
三、Demo以及用预训练模型进行测试
- 下载预训练模型
(这里所说的预训练模型不同于一般意义上迁移学习中的预训练模型,其实是已经在faster rcnn上训练好直接可以用的)
作者只提供了谷歌网盘的地址,为了照顾墙内的同学,我这挂一个vgg16到百度网盘
谷歌网盘(全部模型):下载
百度网盘(vgg16):下载
- 创建文件夹和软链接来使用预训练模型
# 在data文件夹中创建名为voc_2007_trainval的文件夹,把下载并解压好的预训练模型文件放进去
# 以下命令在faster rcnn 根目录下操作
NET=vgg16
TRAIN_IMDB=voc_2007_trainval
mkdir -p output/${NET}/${TRAIN_IMDB}
cd output/${NET}/${TRAIN_IMDB}
ln -s ../../../data/voc_2007_trainval ./default
cd ../../..
- 用demo测试自定义图像
# 同样在根目录下操作(必须这样,不然出错)
python ./tools/demo.py --net vgg16 --dataset pascal_voc
跑完以后大概是这样的效果
想自己加几张图片测试的话,可以把把jpg结尾的图片放进 data/demo里,然后把图片的名称加进 demo.py的代码里
- 使用预训练模型跑测试集
GPU_ID=0
./experiments/scripts/test_faster_rcnn.sh $GPU_ID pascal_voc vgg16
如果报错的话,检查下预训练模型的位置、数据集的位置、软链接这些有没有出问题
跑起来的时候是这个样子滴
跑完以后会得到各个类别上的精度以及平均精度,vgg16的精度大概是0.70+,像我跑出来的结果是比较合理的
如果得到的精度低的离谱,请参考这里
如果只想在CPU上运行,参考这里
四、训练你自己的模型
- 将预训练好的可以用于迁移的模型放在data/imagenet_weights目录下,以vgg16为例如下
mkdir -p data/imagenet_weights # 创建文件夹
cd data/imagenet_weights # 进入文件夹
然后用在该目录下运行这些python代码
import torch
from torch.utils.model_zoo import load_url
from torchvision import models
sd = load_url("https://s3-us-west-2.amazonaws.com/jcjohns-models/vgg16-00b39a1b.pth")
sd['classifier.0.weight'] = sd['classifier.1.weight']
sd['classifier.0.bias'] = sd['classifier.1.bias']
del sd['classifier.1.weight']
del sd['classifier.1.bias']
sd['classifier.3.weight'] = sd['classifier.4.weight']
sd['classifier.3.bias'] = sd['classifier.4.bias']
del sd['classifier.4.weight']
del sd['classifier.4.bias']
torch.save(sd, "vgg16.pth")
- 返回主目录,进行训练
./experiments/scripts/train_faster_rcnn.sh [GPU_ID] [DATASET] [NET]
# GPU_ID 是你希望使用的GPU序号(只有一块GPU就选0)
# NET 是你想使用的网络类型,可选: {vgg16, res50, res101, res152}
# DATASET 是你想要用来训练的数据集,可选:{pascal_voc, pascal_voc_0712, coco}
# 举例:
./experiments/scripts/train_faster_rcnn.sh 0 pascal_voc vgg16
./experiments/scripts/train_faster_rcnn.sh 1 coco res101
注意:如果要训练自己的模型,请删除之前demo时用到的软连接 ln -s ../../../data/voc_2007_trainval ./default
因为新训练好的模型也会保存在这个目录下
- 测试与评估
./experiments/scripts/test_faster_rcnn.sh [GPU_ID] [DATASET] [NET]
# 参数就类似训练的时候,参考上一条
# Examples:
./experiments/scripts/test_faster_rcnn.sh 0 pascal_voc vgg16
./experiments/scripts/test_faster_rcnn.sh 1 coco res101
- 其他
训练好的网络将保存在这个目录 output/[NET]/[DATASET]/default/
测试结果保存在这个目录 output/[NET]/[DATASET]/default/[SNAPSHOT]/
五、在自己的数据上进行训练
最简单的方法,把自己的数据处理成pascal_voc_2007数据的格式,然后覆盖原数据。按照“四”中进行训练
如果制作数据集并替换,之后我会写篇文章专门介绍,在此之前大家可以参考这篇