调用caffemodel中的权重赋值给自己的网络

调用预训练好的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()

你可能感兴趣的:(深度学习,caffe)