【机器学习】截取caffe模型中的某层

  • 2020 年 10 月 19 日
  • AI

一、前言

通常情况下,训练好的caffe模型包含两个文件:

  • prototxt:网络结构描述文件,存储了整个网络的图结构;
  • caffemodel:权重文件,存储了模型权重的相关参数和具体信息

对于某些大型的网络,如果需要对caffe模型中一小部分结构网络做一些性能或精度分析,可以对caffe模型的结构及权重进行截取。caffe源码提供了读取caffe模型的python接口,可以比较方便地进行相应的操作。

二、修改prototxt

以Inception模型为例,如下图左则为使用Netron可视化的模型前3层及prototxt描述的内容,假设要截取第3层,修改后的模型如下图右则所示,修改方法为:

  • 修改输入数据的维度为第2层的输出维度
  • 删除第2层以及第3层后的所有层
  • 修改第3层的bottom值为第一层的top值
「机器学习」截取caffe模型中的某层

三、截取权重数据

首先要下载caffe源码

git clone //github.com/BVLC/caffe

然后使用caffe的python接口读取修改后的prototxt和原始的权重文件caffemodel,接着重新推理,最后保存新的权重文件。

import sys
caffe_root='/your/path/caffe'
sys.path.insert(0, caffe_root + '/python')
import caffe

net = caffe.Net("Inception.prototxt", "Inception.caffemodel", caffe.TRAIN)
res = net.forward()
net.save('Inception_conv2d_2.caffemodel')