使用已有模型初始化caffe网络

import caffe
import numpy as np
import math

# 加载预训练网络
net_ = caffe.Net("deploy.prototxt", "test.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()

猜你喜欢

转载自blog.csdn.net/weixin_38582851/article/details/80579849