train.py 文件源码

def __init__(self):
        if pa.GPU==True:
            caffe.set_device(pa.device)
            caffe.set_mode_gpu()
        else:
            caffe.set_mode_cpu()
        self.solver=caffe.SGDSolver(pa.solver)
        if pa.pretrain!="":
            self.solver.net.copy_from(pa.pretrain)
        self.solver_param=caffe_pb2.SolverParameter()
        with open(pa.solver,'rt') as f:
            pb2.text_format.Merge(f.read(),self.solver_param)

        #self.output_dir=pa.output_dir
        self.solver.net.layers[0].set_queue()