【detectron】pkl模型转.pb模型

detectron训练出来的目标检测模型后缀为.pkl，这种模型在使用的时候必须要有图（graph）以及detectron代码的支持，转为caffe2标准的pb模型后，就可以脱离detectron的代码单独运行，非常方便。（经过翻阅github的相关资料，转为pb模型后可以支持在c++上运行？）

转换所要用到的函数为tools/convert_pkl_to_pb.py。

进入这个函数，主要要填写的部分为，以下标红的部分是必须要修改的，其余是可以按需求修改

1.cfg

    parser.add_argument(
        '--cfg', dest='cfg_file', help='optional config file', default=None,
        type=str)

即cfg的路径。此外cfg文件中的TEST部分添加：

TEST:
  WEIGHTS:#  你的pkl文件路径

2.test_img

test_img是测试转换后的pb文件与原始的pkl文件是否一致，有没有出错。如果添加了test_img，就会进行对比测试，如果没有添加，demo就不会检查，只会输出pb模型文件

    parser.add_argument(
        '--test_img', dest='test_img',
        help='optional test image, used to verify the model conversion',
        default=None,
        type=str)

3.使用设备（device）

    parser.add_argument(
        '--device', dest='device',
        help='Device to run the model on',
        choices=['cpu', 'gpu'],
        default='cpu',
        type=str)

如果你准备在cpu上运行pb文件，那么选择cpy，否则选择gpu

4.pb模型的保存路径

    parser.add_argument(
        '--out_dir', dest='out_dir', help='output dir', default=None,
        type=str)

5.pb模型的名称

    parser.add_argument(
        '--net_name', dest='net_name', help='optional name for the net',
        default="detectron", type=str)

 

修改完上述所有之后，就可以运行啦，会在输出的路径中生成以下几个文件：

• model.pb
• model_init.pb

 

可能出现的问题

1. fpn not support

这个是老版本caffe2和detectron才会出现的情况，新版本的caffe2和detectron就已经做了修复。解决方法是：

（1）从官网上下载新版本的pytorch/modules/detectron，替换老版本的相应位置

（2）从官网上下载新版本的pytorch/caffe2/operators，替换老版本的相应位置

（3）从detectron官网上下载新版本的convert_pkl_to_pb.py，替换老版本。