<commit_before>#!/usr/bin/env python

import google.protobuf
import numpy as np
import matplotlib
from numpy import dtype
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import glob, re, sklearn, collections, argparse, sys, os
import sklearn.metrics
import caffe
import time

'''Script for training a neural net model from gnina grid data.
A model template is provided along with training and test sets of the form
<prefix>[train|test][num].types
Test area, as measured by AUC, is periodically assessed.   At the end graphs are made.
Default is to do dynamic stepping of learning rate, but can explore other methods.
'''
def eval_model(args, trainfile, testfile, reducedtrainfile, reducedtestfile, outname):
    '''run solver for iterations steps, on the given training file,
    every testiter evaluate the roc of bothe the trainfile and the testfile
    return the full predictions for every tested iteration'''
    template = args.model
    testiter = args.test_interval
    iterations = args.iterations
    
    if testiter > iterations: #need to test once
        testiter = iterations
    model = open(template).read().replace('TRAINFILE',trainfile)
    testmodel = model.replace('TESTFILE',testfile)
    trainmodel = model.replace('TESTFILE',trainfile) #for test on train
    if reducedtrainfile != '':        
        reducedtrainmodel = model.replace('TESTFILE', reducedtrainfile)
    else:
        reducedtrainmodel = trainmodel
        reducedtrainfile = trainfile
    if reducedtestfile != '':        
        reducedtestmodel = model.replace('TESTFILE', reducedtestfile)
    else:
        reducedtestmodel = testmodel
        reducedtestfile = testfile
        
    if args.avg_rotations:
        rotations = 24 
        index=testmodel.find(testfile) #add 'rotate = 24' to testmodels if not already there
        endindex=testmodel.find('layer', index)
        rot = testmodel.find("rotate:", index, endindex)
        if rot == -1:
            index = testmodel.find('balanced:', index, endindex)
            index = testmodel.find('\n', index, endindex)
            testmodel = testmodel[:index+1] + '    rotate: %d'% rotations + testmodel[index:]
        
        index=reducedtestmodel.find(reducedtestfile)
        endindex=reducedtestmodel.find('layer', index)
        rot = reducedtestmodel.find("rotate:", index, endindex)
        if rot == -1:
            index = reducedtestmodel.find('balanced:', index, endindex)
            index = reducedtestmodel.find('\n', index, endindex)
            reducedtestmodel = reducedtestmodel[:index+1] + '    rotate: %d' %rotations + reducedtestmodel[index:]
		    

    mode = 'w'
    if args.cont:
        mode = 'a'    
        modelname = '%s_iter_%d.caffemodel' % (outname,args.cont)
        solvername = '%s_iter_%d.solverstate' % (outname,args.cont)
        
    out = open('%s.out' % outname,mode,0)

    pid = os.getpid()
    #very obnoxiously, python interface requires network definition to be in a file
    testproto = 'traintest.%d.prototxt' % pid
    trainproto = 'traintrain.%d.prototxt' % pid
    reducedtestproto = 'trainreducedtest.%d.prototxt' % pid
    reducedtrainproto = 'trainreducedtrain.%d.prototxt' % pid
    with open(testproto,'w') as f:
        f.write(testmodel)    
    with open(trainproto,'w') as f:
        f.write(trainmodel)
    with open(reducedtestproto, 'w') as f:
        f.write(reducedtestmodel) 
    with open(reducedtrainproto, 'w') as f:
        f.write(reducedtrainmodel) 
    solverf = 'solver.%d.prototxt'%pid
    solver_text = '''
    # The train/test net protocol buffer definition
    train_net: "traintest.%d.prototxt"
    test_net: "traintest.%d.prototxt"
    test_net: "traintrain.%d.prototxt"
    test_net: "trainreducedtest.%d.prototxt"
    test_net: "trainreducedtrain.%d.prototxt"
    # The base learning rate, momentum and the weight decay of the network.
    type: "%s"
    base_lr: %f
    momentum: %f
    test_iter: 1
    test_iter: 1
    test_iter: 1
    test_iter: 1
    test_interval: 100000 #we will test manually, these are just here to make caffe happy
    weight_decay: %f
    # The learning rate policy
    lr_policy: "%s"
    gamma: %f
    power: %f
    display: 0
    # reproducible results
    random_seed: %d
    # The maximum number of iterations
    max_iter: %d
    snapshot_prefix: "%s"
    ''' % (pid,pid,pid,pid,pid, args.solver,args.base_lr, args.momentum, args.weight_decay, args.lr_policy, args.gamma, args.power, args.seed, iterations+args.cont,outname)
    with open(solverf,'w') as f:
        f.write(solver_text)
        
    if args.gpu >= 0:
        caffe.set_device(args.gpu)
    caffe.set_mode_gpu()
    
    solver = caffe.get_solver(solverf)
        
    if args.cont:
        solver.restore(solvername)
        solver.testall() #link testnets to train net

    if args.weights:
        solver.net.copy_from(args.weights)
        
    ntests = sum(1 for line in open(reducedtestfile))
    ntrains = sum(1 for line in open(reducedtrainfile))

    testvals = []
    trainvals = []
    bestauc = 0
    bestauci = 0;
    besttestauc = 0
    testauc = 0
    trainauc = 0
    loss = 0
    #print "cnts %d,%d" % (ntrains,ntests)
    for i in xrange(iterations/testiter):
        start = time.time()
        solver.step(testiter)

        print "Train time: %f" % (time.time()-start)

        start = time.time()
        #evaluate test set
        if i == (iterations/testiter)-1 and args.reduced:
            testnet = solver.test_nets[0]
            ntests = sum(1 for line in open(testfile))
        else:
            testnet = solver.test_nets[2]
        y_true = []
        y_score = []
        y_scores = [[] for _ in xrange(ntests)]
        y_affinity = []
        y_predaff = []
        y_predaffs = [[] for _ in xrange(ntests)]
        for x in xrange(ntests):
            res = testnet.forward()
            #MUST copy values out of res as it is return by ref
            y_true.append(float(res['labelout']))
            if 'output' in res:
                y_scores[x].append(float(res['output'][0][1])) 
            else:
                y_scores[x].append(0)
            if 'affout' in res:
                y_affinity.append(float(res['affout']))
                y_predaffs[x].append(float(res['predaff']))
        if args.avg_rotations:
            for _ in xrange(rotations-1):
                print ntests #check if ntests is correct
                for x in xrange(ntests):
                    res = testnet.forward()
                    if 'output' in res:
                        yt = float(res['labelout'])
                        if yt != y_true[x]:
                            print "%dERROR: %f,y_true: %f" %(x,yt, y_true[x]) #sanity check
                        y_scores[x].append(float(res['output'][0][1]))
                    if y_affinity:
                        y_predaffs[x].append(float(res['predaff']))
            #average the scores for the 24 rotations
            for x in xrange(ntests):
                if y_scores:
                    y_score.append(np.mean(y_scores[x]))
                if y_affinity:
                    y_predaff.append(np.mean(y_predaffs[x]))
        else:
            y_score = [row[0] if len(row) > 0 else [] for row in y_scores[0:]]
            y_predaff = [row[0] if len(row) > 0 else [] for row in y_predaffs]
            
        print "Test time: %f" % (time.time()-start)
        if len(np.unique(y_true)) > 1:
            testauc = sklearn.metrics.roc_auc_score(y_true,y_score)
        
        if y_affinity:
            y_predaff = np.array(y_predaff)
            yt = np.array(y_true,np.bool)
            y_affinity = np.array(y_affinity)
            testrmsd = sklearn.metrics.mean_squared_error(y_affinity[yt],y_predaff[yt])
            testvals.append((testauc,y_true,y_score,testrmsd,y_affinity,y_predaff))              
        else:
            testvals.append((testauc,y_true,y_score))
        
        print "Test eval: %f s" % (time.time()-start)
        
        if testauc > besttestauc:
            besttestauc = testauc
            if args.keep_best:
                solver.snapshot() #a bit too much - gigabytes of data
        
        start = time.time()
        #evaluate train set
        start = time.time()
        if i == (iterations/testiter)-1 and args.reduced:
            testnet = solver.test_nets[1]
            ntrains= sum(1 for line in open(trainfile))
        else:
            testnet = solver.test_nets[3]
        y_true = []
        y_score = []
        y_affinity = []
        y_predaff = []
        losses = []
        for x in xrange(ntrains):
            res = testnet.forward()            
            #MUST copy values out of res as it is return by ref
            y_true.append(float(res['labelout']))
            if 'output' in res:
                y_score.append(float(res['output'][0][1]))
            else:
                y_score.append(0)
            if 'loss' in res:
                losses.append(float(res['loss']))
            else:
                losses.append(0)
            if 'affout' in res:
                y_affinity.append(float(res['affout']))
                y_predaff.append(float(res['predaff']))
        
        print "Test train time: %f" % (time.time()-start)
        if len(np.unique(y_true)) > 1:
            trainauc = sklearn.metrics.roc_auc_score(y_true,y_score)            
        loss = np.mean(losses)
        
        if y_affinity:
            y_predaff = np.array(y_predaff)
            y_affinity = np.array(y_affinity)
            trainrmsd = sklearn.metrics.mean_squared_error(y_affinity[yt],y_predaff[yt])                    
            if y_score:
                yt = np.array(y_true,np.bool)
                trainvals.append((trainauc,y_true,y_score,loss,trainrmsd,y_affinity,y_predaff))
            else:
                trainvals.append((0,[],[],0,trainrmsd,y_affinity,y_predaff))
        else:
            trainvals.append((trainauc,y_true,y_score,loss))

        print "Train eval: %f s" % (time.time()-start)
        
        if trainauc > bestauc:
            bestauc = trainauc
            bestauci = i
            
        if args.dynamic: #check for improvement
            lr = solver.get_base_lr()
            if (i-bestauci) > args.step_when: #reduce learning rate
                lr *= args.step_reduce
                solver.set_base_lr(lr)
                bestauci = i #reset 
                bestauc = trainauc #the value too, so we can consider the recovery
            if lr < args.step_end:
                break #end early  
            
        out.write('%.4f %.4f %.6f %.6f'%(testauc,trainauc,loss,solver.get_base_lr()))
        if len(y_affinity):
            out.write(' %.4f %.4f' % (testrmsd,trainrmsd))
        out.write('\n')
        out.flush()
    
    out.close()
    solver.snapshot()
    del solver #free mem
    
    if not args.keep:
        os.remove(solverf)
        os.remove(testproto)
        os.remove(trainproto)
        os.remove(reducedtestproto)
        os.remove(reducedtrainproto)
    return testvals,trainvals


def comma_separated_ints(ints):
     return [int(i) for i in ints.split(',') if i and i != 'None']


def parse_args(argv=None):
    parser = argparse.ArgumentParser(description='Train neural net on .types data.')
    parser.add_argument('-m','--model',type=str,required=True,help="Model template. Must use TRAINFILE and TESTFILE")
    parser.add_argument('-p','--prefix',type=str,required=True,help="Prefix for training/test files: <prefix>[train|test][num].types")
    parser.add_argument('-n','--foldnums',type=comma_separated_ints,required=False,help="Fold numbers to run, default is '0,1,2'",default='0,1,2')
    parser.add_argument('-a','--allfolds',action='store_true',required=False,help="Train and test file with all data folds, <prefix>.types",default=False)
    parser.add_argument('-i','--iterations',type=int,required=False,help="Number of iterations to run,default 10,000",default=10000)
    parser.add_argument('-s','--seed',type=int,help="Random seed, default 42",default=42)
    parser.add_argument('-t','--test_interval',type=int,help="How frequently to test (iterations), default 40",default=40)
    parser.add_argument('-o','--outprefix',type=str,help="Prefix for output files, default <model>.<pid>",default='')
    parser.add_argument('-g','--gpu',type=int,help='Specify GPU to run on',default=-1)
    parser.add_argument('-c','--cont',type=int,help='Continue a previous simulation from the provided iteration (snapshot must exist)',default=0)
    parser.add_argument('-k','--keep',action='store_true',default=False,help="Don't delete prototxt files")
    parser.add_argument('-r', '--reduced', action='store_true',default=False,help="Use a reduced file for model evaluation if exists(<prefix>[_reducedtrain|_reducedtest][num].types)")
    parser.add_argument('--avg_rotations', action='store_true',default=False, help="Use the average of the testfile's 24 rotations in its evaluation results")
    #parser.add_argument('-v,--verbose',action='store_true',default=False,help='Verbose output')
    parser.add_argument('--keep_best',action='store_true',default=False,help='Store snapshots everytime test AUC improves')
    parser.add_argument('--dynamic',action='store_true',default=False,help='Attempt to adjust the base_lr in response to training progress')
    parser.add_argument('--solver',type=str,help="Solver type. Default is SGD",default='SGD')
    parser.add_argument('--lr_policy',type=str,help="Learning policy to use. Default is inv.",default='inv')
    parser.add_argument('--step_reduce',type=float,help="Reduce the learning rate by this factor with dynamic stepping, default 0.5",default='0.5')
    parser.add_argument('--step_end',type=float,help='Terminate training if learning rate gets below this amount',default=0)
    parser.add_argument('--step_when',type=int,help="Perform a dynamic step (reduce base_lr) when training has not improved after this many test iterations, default 10",default=10)
    parser.add_argument('--base_lr',type=float,help='Initial learning rate, default 0.01',default=0.01)
    parser.add_argument('--momentum',type=float,help="Momentum parameters, default 0.9",default=0.9)
    parser.add_argument('--weight_decay',type=float,help="Weight decay, default 0.001",default=0.001)
    parser.add_argument('--gamma',type=float,help="Gamma, default 0.001",default=0.001)
    parser.add_argument('--power',type=float,help="Power, default 1",default=1)
    parser.add_argument('--weights',type=str,help="Set of weights to initialize the model with")
    return parser.parse_args(argv)


if __name__ == '__main__':
    args = parse_args()
    
    #identify all train/test pairs
    pairs = []
    for i in args.foldnums:
        train = '%strain%d.types' % (args.prefix, i)
        test = '%stest%d.types' % (args.prefix, i)
        if not os.path.isfile(train):
            print 'error: %s does not exist' % train
            sys.exit(1)
        if not os.path.isfile(test):
            print 'error: %s does not exist' % test
            sys.exit(1)
        pairs.append((train, test))
    if args.allfolds:
        train = test = '%s.types' % args.prefix
        if not os.path.isfile(train):
            print 'error: %s does not exist' % train
            sys.exit(1)
        pairs.append((train, test))
    
    if len(pairs) == 0:
        print "error: missing train/test files"
        sys.exit(1)
    
    for (train, test) in pairs:
        print train, test
    
    outprefix = args.outprefix
    if outprefix == '':
        outprefix = '%s.%d' % (os.path.splitext(os.path.basename(args.model))[0],os.getpid())
    
    mode = 'w'
    if args.cont:
        mode = 'a'
    
    #train each pair
    testaucs = []
    trainaucs = []
    testrmsds = []
    trainrmsds = []
    alltest = []
    for (train,test) in pairs:
    
        if args.allfolds and train == test:
            crossval = False
            outname = '%s.all' % outprefix
            reducedtrainfile = reducedtestfile = train.replace('.types', '_reduced.types')
        else:
            crossval = True
            m = re.search('%strain(\d+)'%args.prefix,train)
            outname = '%s.%s' % (outprefix,m.group(1))
            reducedtrainfile = train.replace('%strain' % args.prefix,'%s_reducedtrain' % args.prefix)
            reducedtestfile = train.replace('%strain' % args.prefix,'%s_reducedtest' % args.prefix)

        if not os.path.isfile(reducedtrainfile):       
            reducedtrainfile = ''
            print 'error: %s does not exist' % reducedtrainfile
            sys.exit(1)
        if not os.path.isfile(reducedtestfile):       
            reducedtestfile = '' 
            print 'error: %s does not exist' % reducedtestfile
            sys.exit(1)

        test,train = eval_model(args, train, test, reducedtrainfile, reducedtestfile, outname)
        if not crossval:
            continue

        testaucs.append([x[0] for x in test])
        trainaucs.append([x[0] for x in train])
        alltest.append(test)
        if len(test[-1]) > 4:
            testrmsds.append([x[3] for x in test])
            trainrmsds.append([x[3] for x in train])
            
        if np.mean(trainaucs) > 0:
            with open('%s.%s.finaltest' % (outprefix,m.group(1)), mode) as out:
                for (label,score) in zip(test[-1][1],test[-1][2]):
                    out.write('%f %f\n'%(label,score))
                out.write('# AUC %f\n'%test[-1][0])<commit_msg>added DATA_ROOT placeholder and flag to train.py, fixed reduced file not found bug<commit_after>#!/usr/bin/env python

import google.protobuf
import numpy as np
import matplotlib
from numpy import dtype
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import glob, re, sklearn, collections, argparse, sys, os
import sklearn.metrics
import caffe
import time

'''Script for training a neural net model from gnina grid data.
A model template is provided along with training and test sets of the form
<prefix>[train|test][num].types
Test area, as measured by AUC, is periodically assessed.   At the end graphs are made.
Default is to do dynamic stepping of learning rate, but can explore other methods.
'''
def eval_model(args, trainfile, testfile, reducedtrainfile, reducedtestfile, outname):
    '''run solver for iterations steps, on the given training file,
    every testiter evaluate the roc of bothe the trainfile and the testfile
    return the full predictions for every tested iteration'''
    template = args.model
    testiter = args.test_interval
    iterations = args.iterations
    
    if testiter > iterations: #need to test once
        testiter = iterations
    model = open(template).read().replace('TRAINFILE', trainfile).replace('DATA_ROOT', args.data_root)
    testmodel = model.replace('TESTFILE',testfile)
    trainmodel = model.replace('TESTFILE',trainfile) #for test on train
    if reducedtrainfile != '':        
        reducedtrainmodel = model.replace('TESTFILE', reducedtrainfile)
    else:
        reducedtrainmodel = trainmodel
        reducedtrainfile = trainfile
    if reducedtestfile != '':        
        reducedtestmodel = model.replace('TESTFILE', reducedtestfile)
    else:
        reducedtestmodel = testmodel
        reducedtestfile = testfile
        
    if args.avg_rotations:
        rotations = 24 
        index=testmodel.find(testfile) #add 'rotate = 24' to testmodels if not already there
        endindex=testmodel.find('layer', index)
        rot = testmodel.find("rotate:", index, endindex)
        if rot == -1:
            index = testmodel.find('balanced:', index, endindex)
            index = testmodel.find('\n', index, endindex)
            testmodel = testmodel[:index+1] + '    rotate: %d'% rotations + testmodel[index:]
        
        index=reducedtestmodel.find(reducedtestfile)
        endindex=reducedtestmodel.find('layer', index)
        rot = reducedtestmodel.find("rotate:", index, endindex)
        if rot == -1:
            index = reducedtestmodel.find('balanced:', index, endindex)
            index = reducedtestmodel.find('\n', index, endindex)
            reducedtestmodel = reducedtestmodel[:index+1] + '    rotate: %d' %rotations + reducedtestmodel[index:]
		    

    mode = 'w'
    if args.cont:
        mode = 'a'    
        modelname = '%s_iter_%d.caffemodel' % (outname,args.cont)
        solvername = '%s_iter_%d.solverstate' % (outname,args.cont)
        
    out = open('%s.out' % outname,mode,0)

    pid = os.getpid()
    #very obnoxiously, python interface requires network definition to be in a file
    testproto = 'traintest.%d.prototxt' % pid
    trainproto = 'traintrain.%d.prototxt' % pid
    reducedtestproto = 'trainreducedtest.%d.prototxt' % pid
    reducedtrainproto = 'trainreducedtrain.%d.prototxt' % pid
    with open(testproto,'w') as f:
        f.write(testmodel)    
    with open(trainproto,'w') as f:
        f.write(trainmodel)
    with open(reducedtestproto, 'w') as f:
        f.write(reducedtestmodel) 
    with open(reducedtrainproto, 'w') as f:
        f.write(reducedtrainmodel) 
    solverf = 'solver.%d.prototxt'%pid
    solver_text = '''
    # The train/test net protocol buffer definition
    train_net: "traintest.%d.prototxt"
    test_net: "traintest.%d.prototxt"
    test_net: "traintrain.%d.prototxt"
    test_net: "trainreducedtest.%d.prototxt"
    test_net: "trainreducedtrain.%d.prototxt"
    # The base learning rate, momentum and the weight decay of the network.
    type: "%s"
    base_lr: %f
    momentum: %f
    test_iter: 1
    test_iter: 1
    test_iter: 1
    test_iter: 1
    test_interval: 100000 #we will test manually, these are just here to make caffe happy
    weight_decay: %f
    # The learning rate policy
    lr_policy: "%s"
    gamma: %f
    power: %f
    display: 0
    # reproducible results
    random_seed: %d
    # The maximum number of iterations
    max_iter: %d
    snapshot_prefix: "%s"
    ''' % (pid,pid,pid,pid,pid, args.solver,args.base_lr, args.momentum, args.weight_decay, args.lr_policy, args.gamma, args.power, args.seed, iterations+args.cont,outname)
    with open(solverf,'w') as f:
        f.write(solver_text)
        
    if args.gpu >= 0:
        caffe.set_device(args.gpu)
    caffe.set_mode_gpu()
    
    solver = caffe.get_solver(solverf)
        
    if args.cont:
        solver.restore(solvername)
        solver.testall() #link testnets to train net

    if args.weights:
        solver.net.copy_from(args.weights)
        
    ntests = sum(1 for line in open(reducedtestfile))
    ntrains = sum(1 for line in open(reducedtrainfile))

    testvals = []
    trainvals = []
    bestauc = 0
    bestauci = 0;
    besttestauc = 0
    testauc = 0
    trainauc = 0
    loss = 0
    #print "cnts %d,%d" % (ntrains,ntests)
    for i in xrange(iterations/testiter):
        start = time.time()
        solver.step(testiter)

        print "Train time: %f" % (time.time()-start)

        start = time.time()
        #evaluate test set
        if i == (iterations/testiter)-1 and args.reduced:
            testnet = solver.test_nets[0]
            ntests = sum(1 for line in open(testfile))
        else:
            testnet = solver.test_nets[2]
        y_true = []
        y_score = []
        y_scores = [[] for _ in xrange(ntests)]
        y_affinity = []
        y_predaff = []
        y_predaffs = [[] for _ in xrange(ntests)]
        for x in xrange(ntests):
            res = testnet.forward()
            #MUST copy values out of res as it is return by ref
            y_true.append(float(res['labelout']))
            if 'output' in res:
                y_scores[x].append(float(res['output'][0][1])) 
            else:
                y_scores[x].append(0)
            if 'affout' in res:
                y_affinity.append(float(res['affout']))
                y_predaffs[x].append(float(res['predaff']))
        if args.avg_rotations:
            for _ in xrange(rotations-1):
                print ntests #check if ntests is correct
                for x in xrange(ntests):
                    res = testnet.forward()
                    if 'output' in res:
                        yt = float(res['labelout'])
                        if yt != y_true[x]:
                            print "%dERROR: %f,y_true: %f" %(x,yt, y_true[x]) #sanity check
                        y_scores[x].append(float(res['output'][0][1]))
                    if y_affinity:
                        y_predaffs[x].append(float(res['predaff']))
            #average the scores for the 24 rotations
            for x in xrange(ntests):
                if y_scores:
                    y_score.append(np.mean(y_scores[x]))
                if y_affinity:
                    y_predaff.append(np.mean(y_predaffs[x]))
        else:
            y_score = [row[0] if len(row) > 0 else [] for row in y_scores[0:]]
            y_predaff = [row[0] if len(row) > 0 else [] for row in y_predaffs]
            
        print "Test time: %f" % (time.time()-start)
        if len(np.unique(y_true)) > 1:
            testauc = sklearn.metrics.roc_auc_score(y_true,y_score)
        
        if y_affinity:
            y_predaff = np.array(y_predaff)
            yt = np.array(y_true,np.bool)
            y_affinity = np.array(y_affinity)
            testrmsd = sklearn.metrics.mean_squared_error(y_affinity[yt],y_predaff[yt])
            testvals.append((testauc,y_true,y_score,testrmsd,y_affinity,y_predaff))              
        else:
            testvals.append((testauc,y_true,y_score))
        
        print "Test eval: %f s" % (time.time()-start)
        
        if testauc > besttestauc:
            besttestauc = testauc
            if args.keep_best:
                solver.snapshot() #a bit too much - gigabytes of data
        
        start = time.time()
        #evaluate train set
        start = time.time()
        if i == (iterations/testiter)-1 and args.reduced:
            testnet = solver.test_nets[1]
            ntrains= sum(1 for line in open(trainfile))
        else:
            testnet = solver.test_nets[3]
        y_true = []
        y_score = []
        y_affinity = []
        y_predaff = []
        losses = []
        for x in xrange(ntrains):
            res = testnet.forward()            
            #MUST copy values out of res as it is return by ref
            y_true.append(float(res['labelout']))
            if 'output' in res:
                y_score.append(float(res['output'][0][1]))
            else:
                y_score.append(0)
            if 'loss' in res:
                losses.append(float(res['loss']))
            else:
                losses.append(0)
            if 'affout' in res:
                y_affinity.append(float(res['affout']))
                y_predaff.append(float(res['predaff']))
        
        print "Test train time: %f" % (time.time()-start)
        if len(np.unique(y_true)) > 1:
            trainauc = sklearn.metrics.roc_auc_score(y_true,y_score)            
        loss = np.mean(losses)
        
        if y_affinity:
            y_predaff = np.array(y_predaff)
            y_affinity = np.array(y_affinity)
            trainrmsd = sklearn.metrics.mean_squared_error(y_affinity[yt],y_predaff[yt])                    
            if y_score:
                yt = np.array(y_true,np.bool)
                trainvals.append((trainauc,y_true,y_score,loss,trainrmsd,y_affinity,y_predaff))
            else:
                trainvals.append((0,[],[],0,trainrmsd,y_affinity,y_predaff))
        else:
            trainvals.append((trainauc,y_true,y_score,loss))

        print "Train eval: %f s" % (time.time()-start)
        
        if trainauc > bestauc:
            bestauc = trainauc
            bestauci = i
            
        if args.dynamic: #check for improvement
            lr = solver.get_base_lr()
            if (i-bestauci) > args.step_when: #reduce learning rate
                lr *= args.step_reduce
                solver.set_base_lr(lr)
                bestauci = i #reset 
                bestauc = trainauc #the value too, so we can consider the recovery
            if lr < args.step_end:
                break #end early  
            
        out.write('%.4f %.4f %.6f %.6f'%(testauc,trainauc,loss,solver.get_base_lr()))
        if len(y_affinity):
            out.write(' %.4f %.4f' % (testrmsd,trainrmsd))
        out.write('\n')
        out.flush()
    
    out.close()
    solver.snapshot()
    del solver #free mem
    
    if not args.keep:
        os.remove(solverf)
        os.remove(testproto)
        os.remove(trainproto)
        os.remove(reducedtestproto)
        os.remove(reducedtrainproto)
    return testvals,trainvals


def comma_separated_ints(ints):
     return [int(i) for i in ints.split(',') if i and i != 'None']


def parse_args(argv=None):
    parser = argparse.ArgumentParser(description='Train neural net on .types data.')
    parser.add_argument('-m','--model',type=str,required=True,help="Model template. Must use TRAINFILE and TESTFILE")
    parser.add_argument('-p','--prefix',type=str,required=True,help="Prefix for training/test files: <prefix>[train|test][num].types")
    parser.add_argument('-d','--data_root',type=str,required=False,help="Root folder for relative paths in train/test files",default='')
    parser.add_argument('-n','--foldnums',type=comma_separated_ints,required=False,help="Fold numbers to run, default is '0,1,2'",default='0,1,2')
    parser.add_argument('-a','--allfolds',action='store_true',required=False,help="Train and test file with all data folds, <prefix>.types",default=False)
    parser.add_argument('-i','--iterations',type=int,required=False,help="Number of iterations to run,default 10,000",default=10000)
    parser.add_argument('-s','--seed',type=int,help="Random seed, default 42",default=42)
    parser.add_argument('-t','--test_interval',type=int,help="How frequently to test (iterations), default 40",default=40)
    parser.add_argument('-o','--outprefix',type=str,help="Prefix for output files, default <model>.<pid>",default='')
    parser.add_argument('-g','--gpu',type=int,help='Specify GPU to run on',default=-1)
    parser.add_argument('-c','--cont',type=int,help='Continue a previous simulation from the provided iteration (snapshot must exist)',default=0)
    parser.add_argument('-k','--keep',action='store_true',default=False,help="Don't delete prototxt files")
    parser.add_argument('-r', '--reduced', action='store_true',default=False,help="Use a reduced file for model evaluation if exists(<prefix>[_reducedtrain|_reducedtest][num].types)")
    parser.add_argument('--avg_rotations', action='store_true',default=False, help="Use the average of the testfile's 24 rotations in its evaluation results")
    #parser.add_argument('-v,--verbose',action='store_true',default=False,help='Verbose output')
    parser.add_argument('--keep_best',action='store_true',default=False,help='Store snapshots everytime test AUC improves')
    parser.add_argument('--dynamic',action='store_true',default=False,help='Attempt to adjust the base_lr in response to training progress')
    parser.add_argument('--solver',type=str,help="Solver type. Default is SGD",default='SGD')
    parser.add_argument('--lr_policy',type=str,help="Learning policy to use. Default is inv.",default='inv')
    parser.add_argument('--step_reduce',type=float,help="Reduce the learning rate by this factor with dynamic stepping, default 0.5",default='0.5')
    parser.add_argument('--step_end',type=float,help='Terminate training if learning rate gets below this amount',default=0)
    parser.add_argument('--step_when',type=int,help="Perform a dynamic step (reduce base_lr) when training has not improved after this many test iterations, default 10",default=10)
    parser.add_argument('--base_lr',type=float,help='Initial learning rate, default 0.01',default=0.01)
    parser.add_argument('--momentum',type=float,help="Momentum parameters, default 0.9",default=0.9)
    parser.add_argument('--weight_decay',type=float,help="Weight decay, default 0.001",default=0.001)
    parser.add_argument('--gamma',type=float,help="Gamma, default 0.001",default=0.001)
    parser.add_argument('--power',type=float,help="Power, default 1",default=1)
    parser.add_argument('--weights',type=str,help="Set of weights to initialize the model with")
    return parser.parse_args(argv)


if __name__ == '__main__':
    args = parse_args()
    
    #identify all train/test pairs
    pairs = []
    for i in args.foldnums:
        train = '%strain%d.types' % (args.prefix, i)
        test = '%stest%d.types' % (args.prefix, i)
        if not os.path.isfile(train):
            print 'error: %s does not exist' % train
            sys.exit(1)
        if not os.path.isfile(test):
            print 'error: %s does not exist' % test
            sys.exit(1)
        pairs.append((train, test))
    if args.allfolds:
        train = test = '%s.types' % args.prefix
        if not os.path.isfile(train):
            print 'error: %s does not exist' % train
            sys.exit(1)
        pairs.append((train, test))
    
    if len(pairs) == 0:
        print "error: missing train/test files"
        sys.exit(1)
    
    for (train, test) in pairs:
        print train, test
    
    outprefix = args.outprefix
    if outprefix == '':
        outprefix = '%s.%d' % (os.path.splitext(os.path.basename(args.model))[0],os.getpid())
    
    mode = 'w'
    if args.cont:
        mode = 'a'
    
    #train each pair
    testaucs = []
    trainaucs = []
    testrmsds = []
    trainrmsds = []
    alltest = []
    for (train,test) in pairs:
    
        if args.allfolds and train == test:
            crossval = False
            outname = '%s.all' % outprefix
            reducedtrainfile = reducedtestfile = train.replace('.types', '_reduced.types')
        else:
            crossval = True
            m = re.search('%strain(\d+)'%args.prefix,train)
            outname = '%s.%s' % (outprefix,m.group(1))
            reducedtrainfile = train.replace('%strain' % args.prefix,'%s_reducedtrain' % args.prefix)
            reducedtestfile = train.replace('%strain' % args.prefix,'%s_reducedtest' % args.prefix)

        if not os.path.isfile(reducedtrainfile):       
            print 'error: %s does not exist' % reducedtrainfile
            sys.exit(1)
        if not os.path.isfile(reducedtestfile):       
            print 'error: %s does not exist' % reducedtestfile
            sys.exit(1)

        test,train = eval_model(args, train, test, reducedtrainfile, reducedtestfile, outname)
        if not crossval:
            continue

        testaucs.append([x[0] for x in test])
        trainaucs.append([x[0] for x in train])
        alltest.append(test)
        if len(test[-1]) > 4:
            testrmsds.append([x[3] for x in test])
            trainrmsds.append([x[3] for x in train])
            
        if np.mean(trainaucs) > 0:
            with open('%s.%s.finaltest' % (outprefix,m.group(1)), mode) as out:
                for (label,score) in zip(test[-1][1],test[-1][2]):
                    out.write('%f %f\n'%(label,score))
                out.write('# AUC %f\n'%test[-1][0])
{'fraction_non_alphanumeric': 0.12602342506254263, 'fraction_numerical': 0.008642256083693428, 'mean_word_length': 2.446022727272727, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 37, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3445945', 'n_tokens_mistral': 11329, 'n_tokens_neox': 10840, 'n_words': 2509}