caffe——網絡參數轉化

在訓練網絡時能夠利用別人的pre-train model來初始化的網絡,caffe能夠實現兩個網絡參數的轉化,前提條件是轉化的層的參數設計是一致的,如下程序是轉化了三個卷積層和三個全鏈接層的參數,python的代碼以下: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')
相關文章
相關標籤/搜索