From 6f98112788affbd6b6375ad02df66aa596109b71 Mon Sep 17 00:00:00 2001
From: ruri <shipeng1108@163.com>
Date: Mon, 8 Jul 2019 11:20:15 +0800
Subject: [PATCH] check gpu (#2743)

---
 PaddleCV/image_classification/eval.py          |  3 ++-
 PaddleCV/image_classification/infer.py         |  3 ++-
 PaddleCV/image_classification/train.py         |  3 ++-
 PaddleCV/image_classification/utils/utility.py | 18 ++++++++++++++++++
 4 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/PaddleCV/image_classification/eval.py b/PaddleCV/image_classification/eval.py
index 8db943bf..b53b1f38 100644
--- a/PaddleCV/image_classification/eval.py
+++ b/PaddleCV/image_classification/eval.py
@@ -29,7 +29,7 @@ import paddle.fluid as fluid
 import reader_cv2 as reader 
 import models
 from utils.learning_rate import cosine_decay
-from utils.utility import add_arguments, print_arguments
+from utils.utility import add_arguments, print_arguments, check_gpu
 
 parser = argparse.ArgumentParser(description=__doc__)
 add_arg = functools.partial(add_arguments, argparser=parser)
@@ -134,6 +134,7 @@ def eval(args):
 def main():
     args = parser.parse_args()
     print_arguments(args)
+    check_gpu(args.use_gpu)
     eval(args)
 
 
diff --git a/PaddleCV/image_classification/infer.py b/PaddleCV/image_classification/infer.py
index f60073f3..142ed013 100644
--- a/PaddleCV/image_classification/infer.py
+++ b/PaddleCV/image_classification/infer.py
@@ -29,7 +29,7 @@ import paddle.fluid as fluid
 import reader_cv2 as reader
 import models
 import utils
-from utils.utility import add_arguments,print_arguments
+from utils.utility import add_arguments,print_arguments, check_gpu
 
 parser = argparse.ArgumentParser(description=__doc__)
 # yapf: disable
@@ -108,6 +108,7 @@ def infer(args):
 def main():
     args = parser.parse_args()
     print_arguments(args)
+    check_gpu(args.use_gpu)
     infer(args)
 
 
diff --git a/PaddleCV/image_classification/train.py b/PaddleCV/image_classification/train.py
index f6d0f707..bc1eecfa 100755
--- a/PaddleCV/image_classification/train.py
+++ b/PaddleCV/image_classification/train.py
@@ -47,7 +47,7 @@ import reader_cv2 as reader
 import utils
 import models
 from utils.fp16_utils import create_master_params_grads, master_param_to_train_param
-from utils.utility import add_arguments, print_arguments
+from utils.utility import add_arguments, print_arguments, check_gpu
 from utils.learning_rate import cosine_decay_with_warmup
 from dist_train import dist_utils
 
@@ -612,6 +612,7 @@ def train(args):
 def main():
     args = parser.parse_args()
     print_arguments(args)
+    check_gpu(args.use_gpu)
     train(args)
 
 
diff --git a/PaddleCV/image_classification/utils/utility.py b/PaddleCV/image_classification/utils/utility.py
index 33ac1e35..86650ee9 100644
--- a/PaddleCV/image_classification/utils/utility.py
+++ b/PaddleCV/image_classification/utils/utility.py
@@ -60,3 +60,21 @@ def add_arguments(argname, type, default, help, argparser, **kwargs):
         type=type,
         help=help + ' Default: %(default)s.',
         **kwargs)
+
+def check_gpu(use_gpu):
+    """
+    Log error and exit when set use_gpu=true in paddlepaddle
+    cpu version.
+    """
+    err = "Config use_gpu cannot be set as true while you are " \
+          "using paddlepaddle cpu version ! \nPlease try: \n" \
+          "\t1. Install paddlepaddle-gpu to run model on GPU \n" \
+          "\t2. Set use_gpu as false in config file to run " \
+          "model on CPU"
+
+     try:
+        if use_gpu and not fluid.is_compiled_with_cuda():
+            logger.error(err)
+            sys.exit(1)
+    except Exception as e:
+        pass
-- 
GitLab