train.py 文件源码

python
阅读 19 收藏 0 点赞 0 评论 0

项目:QScode 作者: PierreHao 项目源码 文件源码
def __init__(self, solver_prototxt, pretrained_model=None):
        """Initialize the SolverWrapper."""

        self.solver = caffe.SGDSolver(solver_prototxt)
        if pretrained_model is not None:
            print ('Loading pretrained model '
                   'weights from {:s}').format(pretrained_model)
            self.solver.net.copy_from(pretrained_model)      
        self.solver_param = caffe.io.caffe_pb2.SolverParameter()
        with open(solver_prototxt, 'rt') as f:
            text_format.Merge(f.read(), self.solver_param)

        if self.solver_param.solver_mode == 1:
            caffe.set_mode_gpu()
            caffe.set_device(params.gpu_id)
            print 'Use GPU', params.gpu_id, 'to train'
        else:
            print 'Use CPU to train'
        #initial python data layer    
        self.solver.net.layers[0].set_db()
评论列表
文章目录


问题


面经


文章

微信
公众号

扫码关注公众号