caffe——网络参数转化

时间:2022-05-26 23:15:00

在训练网络时可以利用别人的pre-train model来初始化的网络,caffe可以实现两个网络参数的转化,前提条件是转化的层的参数设计是一致的,以下程序是转化了三个卷积层和三个全连接层的参数,python的代码如下:

import caffe
caffe.set_mode_gpu()
train_net = caffe.Net('/home/python_code/caffe/trainmodel.prototxt',
                      '/home/python_code/caffe/gendernet_50000.caffemodel', caffe.TEST)
test_net = caffe.Net('/home/python_code/caffe/deploy.prototxt', caffe.TEST)
test_net.save('/home/python_code/caffe/gendernet.caffemodel')
params = ['conv1', 'conv2', 'conv3', 'fc6', 'fc7', 'fc8']
params_trans = ['conv1', 'conv2', 'conv3', 'fc6', 'fc7', 'fc8']
train_params = {pr: (train_net.params[pr][0].data, train_net.params[pr][1].data) for pr in params}
test_params = {pr: (test_net.params[pr][0].data, test_net.params[pr][1].data) for pr in params_trans}
for pr_train, pr_test in zip(params, params_trans):
    test_params[pr_test][0].flat = train_params[pr_train][0].flat
    test_params[pr_test][1][...] = train_params[pr_train][1]
test_net.save('/home/python_code/caffe/gendernet.caffemodel')