diff --git a/BERT/utils/args.py b/BERT/utils/args.py index 4bb20cc80e62fb6a4482c9542e9d8f899335aeb4..c89c85d716f81f831ea34be841e95377a871052a 100644 --- a/BERT/utils/args.py +++ b/BERT/utils/args.py @@ -20,6 +20,7 @@ from __future__ import print_function import six import argparse +import paddle.fluid as fluid def str2bool(v): # because argparse does not support to parse "true, False" as python diff --git a/ERNIE/run_classifier.py b/ERNIE/run_classifier.py index 409a690f04943ea356509cbba2b2c8af66ec0b43..61fda8f7d48421da90a7f499ad1e72bc2cc792b7 100644 --- a/ERNIE/run_classifier.py +++ b/ERNIE/run_classifier.py @@ -27,7 +27,7 @@ import reader.task_reader as task_reader from model.ernie import ErnieConfig from finetune.classifier import create_model, evaluate from optimization import optimization -from utils.args import print_arguments +from utils.args import print_arguments, check_cuda from utils.init import init_pretraining_params, init_checkpoint from finetune_args import parser @@ -272,5 +272,5 @@ def main(args): if __name__ == '__main__': print_arguments(args) - + check_cuda(args.use_cuda) main(args) diff --git a/ERNIE/run_sequence_labeling.py b/ERNIE/run_sequence_labeling.py index feffb8fe915bce399ab038b8add43b30235430b3..1f752bc622923a6196d77991730d39719b40cac6 100644 --- a/ERNIE/run_sequence_labeling.py +++ b/ERNIE/run_sequence_labeling.py @@ -27,7 +27,7 @@ import reader.task_reader as task_reader from model.ernie import ErnieConfig from optimization import optimization from utils.init import init_pretraining_params, init_checkpoint -from utils.args import print_arguments +from utils.args import print_arguments, check_cuda from finetune.sequence_label import create_model, evaluate from finetune_args import parser @@ -280,4 +280,5 @@ def main(args): if __name__ == '__main__': print_arguments(args) + check_cuda(args.use_cuda) main(args) diff --git a/ERNIE/train.py b/ERNIE/train.py index 4faec96ba07aa6931ba5c96a0cacc53209058924..6cbd7bba15913fded9c0e4e1f921cd3b6cf041f3 100644 --- a/ERNIE/train.py +++ b/ERNIE/train.py @@ -27,7 +27,7 @@ import paddle.fluid as fluid from reader.pretraining import ErnieDataReader from model.ernie import ErnieModel, ErnieConfig from optimization import optimization -from utils.args import print_arguments +from utils.args import print_arguments, check_cuda from utils.init import init_checkpoint, init_pretraining_params from pretrain_args import parser @@ -351,6 +351,7 @@ def train(args): if __name__ == '__main__': print_arguments(args) + check_cuda(args.use_cuda) if args.do_test: test(args) else: diff --git a/ERNIE/utils/args.py b/ERNIE/utils/args.py index b9be634f0f383db61eb667df2345a89262179fd8..ebfb76724243cb40a1ef5ec404ad99a016a19faf 100644 --- a/ERNIE/utils/args.py +++ b/ERNIE/utils/args.py @@ -20,6 +20,8 @@ from __future__ import print_function import six import argparse +import paddle.fluid as fluid + def str2bool(v): # because argparse does not support to parse "true, False" as python @@ -46,3 +48,15 @@ def print_arguments(args): for arg, value in sorted(six.iteritems(vars(args))): print('%s: %s' % (arg, value)) print('------------------------------------------------') + + +def check_cuda(use_cuda, err = \ + "\nYou can not set use_cuda = True in the model because you are using paddlepaddle-cpu.\n \ + Please: 1. Install paddlepaddle-gpu to run your models on GPU or 2. Set use_cuda = False to run models on CPU.\n" + ): + try: + if use_cuda == True and fluid.is_compiled_with_cuda() == False: + print(err) + sys.exit(1) + except Exception as e: + pass