xkrs_manan/RODY
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

训练核心显示修改

Browse Source
This commit is contained in:
JIAKUNHAO 2022-11-18 14:28:10 +08:00
parent 2dac2a6a79
commit 474510b176
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
app/controller/AlgorithmController.py
View File

@ -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},