调用预训练好的caffemodel,将其权重赋值给自己网络中的对应层。基本思路比较简单,三步走:
1、加载预定义的caffemodel;
2、加载要训练的网络;
3、将预训练model的参数赋值给要训练的网络。
import caffe
import numpy as np
import math
# 加载预训练网络
net_ = caffe.Net("deploy.prototxt", "xxx.caffemodel", caffe.TEST)
caffe.set_device(0)
caffe.set_mode_gpu()
# 设置自己要训练的网络
solver = caffe.SGDSolver('solver.prototxt')
# 复制对应层的参数给自己的网络
for param_name in net_.params.keys():
print param_name
if len(net_.params[param_name])==1:
solver.net.params[param_name][0].data[...]=net_.params[param_name][0].data[...]
if len(net_.params[param_name])==2:
solver.net.params[param_name][0].data[...]=net_.params[param_name][0].data[...]
solver.net.params[param_name][1].data[...]=net_.params[param_name][1].data[...]
if len(net_.params[param_name])==3:
solver.net.params[param_name][0].data[...]=net_.params[param_name][0].data[...]
solver.net.params[param_name][1].data[...]=net_.params[param_name][1].data[...]
solver.net.params[param_name][2].data[...]=net_.params[param_name][2].data[...]
# 开始训练
solver.solve()