main.py 文件源码

def main(argv):
    pprint.pprint(tf.app.flags.FLAGS.__flags)

    flags = tf.app.flags.FLAGS        

    graph = tf.Graph()
    os.environ["CUDA_VISIBLE_DEVICES"]=str(flags.gpu)
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.22, allow_growth=True)
    with graph.as_default():
        with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options),graph=graph) as session:
            all_tests = test_list.tests
            tasks = task_list.tasks

            if flags.train == 'optimizer':
                train_optimizer(tasks[flags.task])
            elif flags.train == 'optimizee':
                train_optimizee(all_tests[flags.task])
            elif flags.train == 'optimizer_train_optimizee':
                optimizer_train_optimizee(tasks[flags.task])
            elif flags.train == 'test':
                test(tasks[flags.task])