caffe利用snapshot从断点恢复训练

约定不等于承诺〃 2022-07-14 13:18 132阅读 0赞

训练网络时迭代次数往往需要数万次，需要训练几天，如果突发什么意外（断电）训练停止了岂不要从头训练。

其实借用snapshot机制，比如每隔1万次迭代保存一下网络参数，然后下次训练判断有无snapshot，有的话从snapshot恢复参数就可以了。

下面是一个例子

#coding:utf-8
    import caffe
    import numpy as np
    import os
    
    from caffe.proto import caffe_pb2
    import google.protobuf as pb2
    
    maxiteration=0
    
    class SolverWrapper(object):
        """A simple wrapper around Caffe's solver.
        This wrapper gives us control over he snapshotting process, which we
        use to unnormalize the learned bounding-box regression weights.
        """
    
        def __init__(self, solver_prototxt, output_dir,
                     pretrained_model=None):
            """Initialize the SolverWrapper."""
            self.output_dir = output_dir
    
            self.solver = caffe.SGDSolver(solver_prototxt)
    #判断输出路径是否有snapshot文件
            flist=os.listdir('/home/hj/py-R-FCN/output/rfcn_end2end_ohem/voc_0712_trainval')
            global maxiteration
            if len(flist)>0:
                maxit=0
                maxid=0
    #得到最新的snapshot文件
                for i in range(len(flist)):
                    ite=int(flist[i].split('_')[-1].split('.')[0])
                    if ite>maxit:
                        maxit=ite
                        maxid=i
                print 'resume training from iteration ' ,maxit
                maxiteration=maxit
                print 'maxiteration:',maxiteration
    #拷贝参数到网络
                self.solver.net.copy_from(os.path.join('/home/hj/py-R-FCN/output/rfcn_end2end_ohem/voc_0712_trainval',flist[maxid]))
    #如果没有snapshot文件则用imagenet上训练得到的模型初始化参数
            elif 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_pb2.SolverParameter()
            with open(solver_prototxt, 'rt') as f:
                pb2.text_format.Merge(f.read(), self.solver_param)
            ...
            ...
            ...
    
    #定义snapshot函数
        def snapshot(self):
            """Take a snapshot of the network after unnormalizing the learned
            bounding-box regression weights. This enables easy use at test-time.
            """
            net = self.solver.net
    
    
            infix = ('_' + cfg.TRAIN.SNAPSHOT_INFIX
                     if cfg.TRAIN.SNAPSHOT_INFIX != '' else '')
    #得到snapshot的文件名
            filename = (self.solver_param.snapshot_prefix + infix +
                        '_iter_{:d}'.format(self.solver.iter+maxiteration) + '.caffemodel')
    #加上路径
            filename = os.path.join(self.output_dir, filename)
            net.save(str(filename))
            print 'Wrote snapshot to: {:s}'.format(filename)

snapshot函数在你训练时可以每隔1万次运行一次

[caffe 在已有模型上继续训练][caffe]

[caffe]: http://blog.csdn.net/u014114990/article/details/47781233