调用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()
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值