def main(_): if args.use_gpu: print("[*] GPU\n") os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_idx # gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_mem) # with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess: with tf.Session(config=config) as sess: model = lowlight_enhance(sess) if args.phase == 'train': lowlight_train(model) elif args.phase == 'test': lowlight_test(model) else: print('[!] Unknown phase') exit(0) else: print("[*] CPU\n") with tf.Session() as sess: model = lowlight_enhance(sess) if args.phase == 'train': lowlight_train(model) elif args.phase == 'test': lowlight_test(model) else: print('[!] Unknown phase') exit(0)
def main(input_file, use_gpu=False): print('called main') if use_gpu: print("[*] GPU\n") os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_idx gpu_options = tf.GPUOptions( per_process_gpu_memory_fraction=args.gpu_mem) with tf.Session(config=tf.ConfigProto( gpu_options=gpu_options)) as sess: model = lowlight_enhance(sess) lowlight_test(input_file, model) else: print("[*] CPU\n") with tf.Session() as sess: model = lowlight_enhance(sess) lowlight_test(input_file, model) #if __name__ == '__main__': # main(input_file, use_gpu=False) #tf.app.run()