【機器學習】截取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