deeplab--trainer.py

来源:互联网 发布:公司日报管理系统源码 编辑:程序博客网 时间:2024/06/06 07:36

感觉作者写的run.py实在是太绕了。 写了好几个模块,最终的目的就是配置各种变量。可能作者想在训练各种数据集的时候直接修改 run.py里面的几个参数就可以。其实直接运行一条caffe train …就可以了。anyway, 跟着作者的脚本学了一下python.

import osimport subprocessimport shutilfrom tools import model_finder, file_editor#train 变量设置#打印 环境变量,train为训练类别设置,init silver #os.environ:环境变量设置:#[os](http://kuanghy.github.io/2015/08/02/python-os)×××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××def train_variables(type_):    print os.environ['MODEL']    #打印model的环境变量,这里的mdoel是在 run.py里面设置的    train = 'train' if type_==1 else 'trainval'    #train方式设置 默认train    init = '/init_' + os.environ['MODEL'] + '.caffemodel' if type_==1 else '/init2_' + os.environ['MODEL'] + '.caffemodel'    #比如MODEL='DEEPLABV2-VGG16' 那么在相应的model文件夹下就会找init_DEEPLABV2-VGG16.caffemodel    solver = 'solver' if type_==1 else 'solver2'        #solver变量设置    train_set=train + os.environ['train_set_SUFFIX']   #设置train和是否打乱    return train, init, solver, train_set#train_variables()函数是设置train的参数×××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××def train_text_maker(train_set):    file1 = os.environ['LIST_DIR'] + "/" + train_set + ".txt"    file2 = os.environ['LIST_DIR'] + "/" + os.environ['train_set_STRONG'] + ".txt"    if int(os.environ['train_set_WEAK_LEN']) ==0:        train_set_WEAK=train_set + '_diff_' + os.environ['train_set_STRONG']        file_output = os.environ['LIST_DIR'] + "/" + train_set_WEAK + ".txt"        if not os.path.isfile(file_output):            command = 'comm -3 {0} {1} > {2}'.format(file1, file2, file_output)            subprocess.call(command, shell=True)    else:        train_set_WEAK= train_set + '_diff_' + os.environ['train_set_STRONG'] + '_head' + os.environ['train_set_WEAK_LEN']               file3 = os.environ['train_set_WEAK_LEN']        file_output = os.environ['LIST_DIR'] + "/" + train_set_WEAK + ".txt"        if not os.path.isfile(file_output):            command = 'comm -3 {0} {1} | head -n {2} > {3}'.format(file1, file2, file3, file_output)            subprocess.call(command, shell=True)×××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××def train_prototxt_maker(train, init, solver, train_set):    model=os.environ['EXP'] + '/model/' +os.environ['NET_ID'] + init #change this    if not os.path.isfile(model): model=model_finder(os.environ['EXP']+ '/model/' + os.environ['NET_ID'])    for variable in ['train', solver]:        file1= os.environ['CONFIG_DIR'] + '/' + variable + '.prototxt'        file_output = os.environ['CONFIG_DIR'] + '/' + variable + '_' + train_set + '.prototxt'        shutil.copyfile(file1, file_output)        file_editor(file_output, train_set=train_set, test_set='')    return model×××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××def train_runner(solver, train_set, model):    cmd = os.environ['CAFFE_DIR'] + os.environ['CAFFE_BIN'] + ' train' \    ' --solver=' + os.environ['CONFIG_DIR'] + '/' + solver + '_' + train_set + '.prototxt' \+ ' --weights=' + model + ' --gpu=' + os.environ['DEV_ID'] #change solver  这里就是用编译好的caffe train  solver  weights  gpu    所以运行的时候可以直接运行这行代码,把选供应的参数对应好就可以    print 'Running ' + cmd #运行train 进行训练    subprocess.call(cmd, shell=True)×××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××def trainer(type_=1):    train, init, solver, train_set = train_variables(type_)    train_text_maker(train_set)    model = train_prototxt_maker(train, init, solver, train_set)    print 'Training' + str(type_) + ' net ' + os.environ['EXP'] + '/' + os.environ['NET_ID']    train_runner(solver, train_set, model) ×××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××××
原创粉丝点击