From 474510b1767535549bc7af653434fd15ae666515 Mon Sep 17 00:00:00 2001 From: JIAKUNHAO Date: Fri, 18 Nov 2022 14:28:10 +0800 Subject: [PATCH] =?UTF-8?q?=E8=AE=AD=E7=BB=83=E6=A0=B8=E5=BF=83=E6=98=BE?= =?UTF-8?q?=E7=A4=BA=E4=BF=AE=E6=94=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/controller/AlgorithmController.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/app/controller/AlgorithmController.py b/app/controller/AlgorithmController.py index 647bc68..cc79d94 100644 --- a/app/controller/AlgorithmController.py +++ b/app/controller/AlgorithmController.py @@ -372,17 +372,17 @@ def Export_model_RODY(params_str): @obtain_train_param() def returnTrainParams(): - # nvmlInit() - # gpuDeviceCount = nvmlDeviceGetCount() # 获取Nvidia GPU块数 - # _kernel = [f"cuda:{a}" for a in range(gpuDeviceCount)] + nvmlInit() + gpuDeviceCount = nvmlDeviceGetCount() # 获取Nvidia GPU块数 + _kernel = [f"cuda:{a}" for a in range(gpuDeviceCount)] params_list = [ {"index": 0, "name": "epochnum", "value": 10, "description": '训练轮次', "default": 100, "type": "I", 'show': True}, {"index": 1, "name": "batch_size", "value": 4, "description": '批次图像数量', "default": 1, "type": "I", 'show': True}, {"index": 2, "name": "img_size", "value": 640, "description": '训练图像大小', "default": 640, "type": "I", 'show': True}, - {"index": 3, "name": "device", "value": "0", "description": '训练核心', "default": "cuda", "type": "S", - "items": '', 'show': True}, # _kernel + {"index": 3, "name": "device", "value": 'CUDA', "description": '训练核心', "default": 'CUDA', "type": "E", + "items": _kernel, 'show': True}, # _kernel {"index": 4, "name": "saveModDir", "value": "E:/alg_demo-master/alg_demo/app/yolov5/best.pt", "description": '保存模型路径', "default": "./app/maskrcnn/saved_model/test.pt", "type": "S", 'show': False},