【机器学习】截取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')

VirMach 便宜 VPS

QNews

QNews