trainCNN.py 文件源码

def train(solver_prototxt_filename, init, init_type):
      for line in fileinput.input('train_net.sh', inplace=True):
              if '-solver' in line:
                 tmp = line.split('-solver')
                 if init==None:
                     print tmp[0]+" -solver "+ solver_prototxt_filename
                 elif init_type == 'fin':
                     print tmp[0]+" -solver "+ solver_prototxt_filename +" -weights " + init # .caffemodel file requiered for finetuning
                 elif init_type == 'res':
                     print tmp[0]+" -solver "+ solver_prototxt_filename +" -snapshot " + init # .solverstate file requiered for resuming training
                 else:
                     raise ValueError("No specific init_type defined for pre-trained network "+init)
              else:
                     print line,
      os.system("chmod +x train_net.sh")
      os.system('./train_net.sh')