[pycaffe] fix CPU / GPU switch in example scripts

shelhamer · shelhamer · commit 3f68c51753cd · 2015-03-08T00:16:26.000-08:00
diff --git a/python/caffe/detector.py b/python/caffe/detector.py
@@ -24,7 +24,7 @@ class Detector(caffe.Net):
     Detector extends Net for windowed detection by a list of crops or
     selective search proposals.
     """
-    def __init__(self, model_file, pretrained_file, gpu=False, mean=None,
+    def __init__(self, model_file, pretrained_file, mean=None,
                  input_scale=None, raw_scale=None, channel_swap=None,
                  context_pad=None):
         """
diff --git a/python/classify.py b/python/classify.py
@@ -96,17 +96,19 @@ def main(argv):
     if args.channel_swap:
         channel_swap = [int(s) for s in args.channel_swap.split(',')]
 
-    # Make classifier.
-    classifier = caffe.Classifier(args.model_def, args.pretrained_model,
-            image_dims=image_dims, gpu=args.gpu, mean=mean,
-            input_scale=args.input_scale, raw_scale=args.raw_scale,
-            channel_swap=channel_swap)
-
     if args.gpu:
+        caffe.set_mode_gpu()
         print("GPU mode")
     else:
+        caffe.set_mode_cpu()
         print("CPU mode")
 
+    # Make classifier.
+    classifier = caffe.Classifier(args.model_def, args.pretrained_model,
+            image_dims=image_dims, mean=mean,
+            input_scale=args.input_scale, raw_scale=args.raw_scale,
+            channel_swap=channel_swap)
+
     # Load numpy array (.npy), directory glob (*.jpg), or image file.
     args.input_file = os.path.expanduser(args.input_file)
     if args.input_file.endswith('npy'):
diff --git a/python/detect.py b/python/detect.py
@@ -107,16 +107,19 @@ def main(argv):
     if args.channel_swap:
         channel_swap = [int(s) for s in args.channel_swap.split(',')]
 
+    if args.gpu:
+        caffe.set_mode_gpu()
+        print("GPU mode")
+    else:
+        caffe.set_mode_cpu()
+        print("CPU mode")
+
     # Make detector.
-    detector = caffe.Detector(args.model_def, args.pretrained_model,
-            gpu=args.gpu, mean=mean,
+    detector = caffe.Detector(args.model_def, args.pretrained_model, mean=mean,
             input_scale=args.input_scale, raw_scale=args.raw_scale,
             channel_swap=channel_swap,
             context_pad=args.context_pad)
 
-    if args.gpu:
-        print('GPU mode')
-
     # Load input.
     t = time.time()
     print("Loading input...")
