From cc646f2807fb728cf44971d6794538ac0b1569e7 Mon Sep 17 00:00:00 2001 From: ykkk2333 Date: Tue, 1 Mar 2022 06:26:31 +0000 Subject: [PATCH] add use xpu option,*test=kunlun --- train.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index 8b01bef614..35526683c6 100644 --- a/train.py +++ b/train.py @@ -120,6 +120,12 @@ def parse_args(): help='The option of train profiler. If profiler_options is not None, the train ' \ 'profiler is enabled. Refer to the paddleseg/utils/train_profiler.py for details.' ) + parser.add_argument( + '--device', + dest='device', + help='Device place to be set, which can be GPU, XPU, CPU', + default='gpu', + type=str) return parser.parse_args() @@ -137,8 +143,13 @@ def main(args): ['-' * 48]) logger.info(info) - place = 'gpu' if env_info['Paddle compiled with cuda'] and env_info[ - 'GPUs used'] else 'cpu' + if args.device == 'gpu' and env_info[ + 'Paddle compiled with cuda'] and env_info['GPUs used']: + place = 'gpu' + elif args.device == 'xpu' and paddle.is_compiled_with_xpu(): + place = 'xpu' + else: + place = 'cpu' paddle.set_device(place) if not args.cfg: