从.weights文件中提取权重数据
前言
我在做有关YOLO的工作,需要从训练好的权重文件中将数据提取出来另作他用,因此记录下提取数据的过程。
.weights文件
.weights文件是Darknet所保存的文件,内含训练出来的权重数据,存储格式特殊,是不能直接用Notepad等工具直接打开或通过python库读出权重数值来的。因此我们需要将其转换为tensorflow的文件(.ckpt文件)或.h5文件来提取权重数据。我做的是将其转换为.h5文件。
将.weights文件转换成.h5
转换脚本convert.py
转换脚本:convert.py,下载地址如下:
https://github.com/qqwweee/keras-yolo3
脚本会根据.weights文件每一层的起始来进行相应转化,有时会报出Unsupported section header type: {}的错误,这是因为你这一层的名字没有包含在分类里。比如我的region_0这一层报错,那么就需要在代码段中添加:
其他准备工作
除了准备好上述的convert.py之外,还需要.cfg文件,可以到github上找到相应版本的(与你要提取的.weights文件版本一致,注意有没有tiny的区别)。
转换成.h5
从cmd中cd到当前文件夹(.weights文件、.cfg文件、convert.py都放在当前文件夹),利用python convert.py XXX.cfg XXX.weights YYY.h5即可生成.h5文件。
从.h5文件提取数据
关于.h5文件
.h5文件长得有点像文件夹,通过一层一层的树状结构来存储数据,利用python自带的h5py库可以将其中数据导入出来。
摸清.h5的子文件夹名字
首先要搞清楚.h5文件下的第一层文件的名字,以我的.h5文件为例:
import h5py
import numpy as np
f = h5py.File('yolov2_tiny.h5','r')
for key in f:
print("key : " + key)
然后就可以在输出处得到第一层的根文件夹名:
答案是只有一个model_weights,那么model_weights下边又有哪些子文件夹呢?
import h5py
import numpy as np
f = h5py.File('yolov2_tiny.h5','r')
for key in f['model_weights']:
print("key : " + key)
输出结果是:
以此类推,终于搞清楚第一个权重的存储位置是/model_weights/conv2d_1/conv2d_1/kernel:0,因此:
import h5py
import numpy as np
f = h5py.File('yolov2_tiny.h5','r')
#for key in f['model_weights']:
# print("key : " + key)
#for group in f['model_weights']['conv2d_2']['conv2d_2']:
# print("group : " + group)
print(f['model_weights']['conv2d_1']['conv2d_1']['kernel:0'][:])
可得到第一个3 × 3 × 3 × 16的卷积核数据,结果如下:
不过输出的顺序可能并不是如我们所愿的顺序,后续还需要用python或者matlab来重新排序(我感觉还是挺繁的,不知有无好办法)。