跑通Pytorch1.0版本的Faster-RCNN，并训练自己的数据集（一）

我的环境： Ubuntu 18.04 + cuda 10.1 + pytorch=1.0.0 + torchvision=0.2.1 + python=3.6.2

 

代码参考的是github上大神的开源代码https://github.com/jwyang/faster-rcnn.pytorch

 

Step一、下载安装

1、首先代码下载，自己down或者git命令clone都可，需要注意的是repo的master分支是pytroch=0.4版本的，所以需要切换到pytorch-1.0分支

git命令

git clone https://github.com/jwyang/faster-rcnn.pytorch.git

在faster-rcnn.pytorch文件夹目录下终端运行

git checkout pytorch-1.0

#打开terminal cd 到faster-rcnn.pytorch文件，同时在内创建一个data文件夹
cd faster-rcnn.pytorch && mkdir data
#并在data文件夹下创建一个pretrained_model，用来存放之后的预训练模型
cd data && mkdir pretrained_model
#安装依赖的pyhon包，这一步若是报错说没有权限访问安装就在开头加上 sudo
#注意scipy版本，不需要太高，我的scipy=1.2.1
pip install -r requirements.txt

接着回退到根目录下运行

cd ..
cd lib
python setup.py build develop

2、然后是预训练模型的下载

作者给出了两种特征提取网络的Pretrained Models，下载链接：

• VGG 16：VT Server
• ResNet 101 VT Server

下载完成后把他们放到data/pretrained_model目录下。

3、数据集的准备

这边使用的是Pascal_VOC数据集

下载数据集并创建软连接（便于多个项目使用同一数据集）

下载训练、验证测试数据集training, validation, test data 和 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

将三个文件解压到同一个文件夹内VOCdevkit

tar xvf VOCtrainval_06-Nov-2007.tar
tar xvf VOCtest_06-Nov-2007.tar
tar xvf VOCdevkit_08-Jun-2007.tar

最终在VOCdevkit目录下有VOC2007/Annotations等文件夹，数据集格式如图所示

VOCdevkit
└── VOC2007
    ├── Annotations
    ├── ImageSets
    │   └── Main
    │       ├── test.txt
    │       ├── train.txt
    │       ├── trainval.txt
    │       └── val.txt
    └── JPEGImages
#VOC数据集的文件夹tree结构：其中Annotations内放xml标注文件，JPEGImages内放图片，
#ImageSets/Main/内的四个txt文件分别是测试集、训练集、训练验证集、验证集。

cd 先前创建的data文件夹，创建软连接 ln -s $VOCdevkit \home\VOCdevkit2007

删除软连接方式：在软连接所在路径下 rm -rf VOCdevkit2007
 

Step二、训练

CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_net.py \
                   --dataset pascal_voc --net vgg16 \
                   --bs $BATCH_SIZE --nw $WORKER_NUMBER \
                   --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \
                   --cuda

设定参数

CUDA_VISIBLE_DEVICES=0 python trainval_net.py --dataset pascal_voc --net res101 --epochs 20 --bs 1 --nw 4 --lr 1e-2 --lr_decay_step 8 --use_tfb --mGPUs --cuda

参数的意义：

    CUDA_VISIBLE_DEVICES：使用哪一个GPU
    –dataset pascal_voc：用pascal数据集

 test_net.py --dataset pascal_voc --net vgg16 --checksession 1 --checkepoch 6 --checkpoint 416 --cuda

    –net res101：用res101作为特征提取网络，可选vgg16
    –epochs 20：20个epoch，每个epoch都过一遍全部的训练图。
    –bs 1：batch_size=1
    –nw 4：num_works=4，4线程进行读取，我的显卡显存是11G，如果显存小的话就设为1
    –lr 1e-2： 初始学习率是0.01；可选参数
    –lr_decay_step 8： 每几个epoch学习率衰减一次（默认衰减一次0.1，通过decay_gamma可调）；可选参数
    –use_tfb： 使用tensorboardX实现记录和可视化，不用就不写；可选参数
    –mGPUs： 多GPU训练，不用就不写该命令；可选参数
    –cuda：使用cuda；

 

Step三、测试

 

测试训练完成模型的效果

python test_net.py --dataset pascal_voc --net vgg16 --checksession $SESSION --checkepoch $EPOCH --checkpoint $CHECKPOINT --cuda

e.g.  SESSION=1, EPOCH=6, CHECKPOINT=416

python test_net.py --dataset pascal_voc --net vgg16 --checksession 1 --checkepoch 6 --checkpoint 416 --cuda

批量测试图片

run demo.py

python demo.py --dataset pascal_voc --net res101  --cfg cfgs/res101.yml --load_dir models --checksession 1 --checkepoch 28 --checkpoint 2813 --image_dir images --cuda

参数的意义：

    –dataset pascal_voc：用pascal数据集
    –net res101：用res101作为特征提取网络，可选vgg16
    –cfg cfgs/res101.yml：加载res101的配置
    –load_dir models --checksession 1 --checkepoch 28 --checkpoint 2813：这四个参数共同决定读取models/文件夹中训练好的faster_rcnn_1_28_2813.pth权重文件
    –image_dir path/to/your/image：需要批量检测的图像文件路径
    –bs 1：batch_size=1
    –cuda：使用cuda
    测试好的图片输出在与image相同的文件夹内，其命名为原图名字+_det。
 

下一篇介绍如何训练测试自己的数据集