From 222f2c35511223bb0d21852616fccc571587359b Mon Sep 17 00:00:00 2001 From: morotyo Date: Sat, 1 Oct 2016 15:44:16 +0900 Subject: [PATCH] Update optimize_image.py for use GPU --- optimize_image.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/optimize_image.py b/optimize_image.py index a555d9e02..325e01bc0 100755 --- a/optimize_image.py +++ b/optimize_image.py @@ -177,6 +177,12 @@ def main(): # Load network sys.path.insert(0, os.path.join(args.caffe_root, 'python')) import caffe + if settings.caffevis_mode_gpu: + caffe.set_mode_gpu() + print 'CaffeVisApp mode (in main thread): GPU' + else: + caffe.set_mode_cpu() + print 'CaffeVisApp mode (in main thread): CPU' net = caffe.Classifier( args.deploy_proto, args.net_weights,