自定义模型规则 ============== 1、数据准备 ----------- ​平台支持8种任务类型的自定义模型训练，包括目标检测、语义分割、实例分割、图像分类、文字检测、文字识别、旋转目标检测、其他。 自定义模型使用平台提供的数据集进行训练。不同任务类型的数据集，拆分格式不同，详情请见 :doc:`图像数据拆分格式 <./dataset/dataset_split_format>`。 用户选择任务类型：其他，创建新的数据集，平台只解析zip包中的图片数据，不解析相关的标签数据，最终以zip包中，实际的文件夹层级存储。且该数据集仅支持自定义算法训练。 注：若zip包中，第一层只有一个文件夹，则将第二层作为根路径。原zip包目录结构示例： :: ├── data ├── annotations ├── train.json ├── val.json ├── test.json ├── images ├── train ├── example1.jpg ├── val ├── example2.jpg ├── test ├── example3.jpg 上传平台后，数据存储的目录结构如下所示： :: ├── annotations ├── train.json ├── val.json ├── test.json ├── images ├── train ├── example1.jpg ├── val ├── example2.jpg ├── test ├── example3.jpg 2、模型准备 ----------- 2.1 目录结构 ~~~~~~~~~~~~ 自定义模型压缩成zip上传平台，zip中的目录结构包括： :: ├── train.py ├── predict.py ├── requirements.txt ├── weights ├── 其他文件或文件夹 其中： - train.py 为训练脚本文件 - predict.py 为预测脚本文件 - requirements.txt 为环境配置文件 - weights 为权重文件存放目录 注意：自定义模型目录结构中，train.py、predict.py 和 requirements.txt 文件的名称不允许修改。权重文件目录weights必须存在。 2.2 环境配置文件 ~~~~~~~~~~~~~~~~ 2.2.1 基础依赖 ^^^^^^^^^^^^^^ 平台默认的安装依赖的命令为： :: pip3 install -U -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple 平台支持两种格式的 requirements.txt 写法： 1. 直接将相关依赖项写入 requirements.txt 文件中，文件内容如下所示： :: matplotlib>=3.3 numpy>=1.23.5 opencv-python>=4.1.1 pillow>=10.3.0 ... 2. 模型安装所需依赖项写在其他文件中，requirements.txt 中引用相关文件，文件内容如下所示： :: -r requirements/base.txt -r requirements/production.txt 此时，自定义模型zip包中的目录结构，需包含： :: ├── train.py ├── requirements.txt ├── requirements（文件夹） ├── base.txt ├── production.txt 2.2.1 其他依赖 ^^^^^^^^^^^^^^ 无法通过pip安装的额外依赖，在项目中，将这些依赖的内容写入sh文件中，文件名：install_dependencies.sh。 （注：此文件名称不允许修改。） 示例：通过mim命令，安装mmengine、mmcv的脚本。 :: #!/bin/bash # 定义 mim 的安装命令 python_path=$1 PIP=${python_path}/bin/pip MIM=${python_path}/bin/mim # 安装mim $PIP install openmim -i https://pypi.tuna.tsinghua.edu.cn/simple # 安装 torch、torchvision $PIP install torch==2.1.0 torchvision==0.16.0 -i https://pypi.tuna.tsinghua.edu.cn/simple # 安装 mmengine $MIM install mmengine -i https://pypi.tuna.tsinghua.edu.cn/simple # 安装 mmcv $MIM install "mmcv>=2.0.0rc4,<2.2.0" -i https://pypi.tuna.tsinghua.edu.cn/simple 注： - python_path，是平台虚拟环境的路径。需要在sh脚本中，允许接收该参数。 - install_dependencies.sh 文件名称，不允许修改 - sh脚本文件中若出现以下非法命令，则会提示脚本非法。 :: 'rm', 'kill', 'remove', 'mv', 'reboot', 'shutdown', 'gcc', 'g++', 'clang++', 'cmake', 'make', 'chown', 'chmod', 'chgrp', 'source', 'crontab', 'eval', 'exec', 'set -e', 'mkfs', 'fdisk' 2.3 训练脚本 ~~~~~~~~~~~~ - 训练脚本模板 训练脚本模板中包括：参数解析方法、训练方法等，用户可基于此模版编写自己的训练脚本。 :: import argparse import os import sys import time from pathlib import Path FILE = Path(__file__).resolve() ROOT = FILE.parents[0] # root directory if str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # add ROOT to PATH ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative def train(args): """ 训练逻辑 """ pass def parse_args(known=False): """ 训练脚本参数 """ parser = argparse.ArgumentParser() parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path") parser.add_argument("--work-dir", default=ROOT / "runs/train", help="save to project/name") parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") parser.add_argument("--label-class", type=str, default="", help="label class") parser.add_argument("--label-class-count", type=int, default=1, help="label class count") parser.add_argument("--epochs", type=int, default=10, help="total training epochs") parser.add_argument("--batch-size", type=int, default=32, help="total batch size for all GPUs") args = parser.parse_args() return args if __name__ == "__main__": args = parse_args() train(args) - 参数解析/全局变量替换 平台提供两种方式，对自定义模型训练脚本 train.py 中所需的参数进行赋值。 方法一：通过args参数赋值： 上传自定义模型时，平台会自动解析训练脚本 train.py 中的参数。在自定义模型列表—训练参数配置页面，创建训练参数配置，平台直接将解析到的训练参数回显到弹窗中（训练参数）。参数的名称与训练脚本中接收的参数一致，用户只需修改参数的值即可保存为参数配置模版。 注: 对于平台直接赋值的参数，前端页面不显示。实际训练时，平台会自动给参数传递实际的值。 方法二：通过全局变量赋值： 训练脚本 train.py 中，若使用了平台支持的全局变量，在实际训练时，平台会自动将脚本中的全局变量替换为实际的值。 args参数及全局变量示例，如下表所示。 .. list-table:: :widths: 20 20 15 40 20 :header-rows: 1 * - args参数名称 - 全局变量名称 - 参数类型 - 参数说明 - 赋值方式 * - work-dir - WORK_DIR_PATH - string - 训练结果保存路径 - 平台赋值 * - data - DATA_PATH - string - 数据集路径 - 平台赋值 * - train-data - TRAIN_DATA_PATH - string - 训练集路径 - 平台赋值 * - val-data - VAL_DATA_PATH - string - 验证集路径 - 平台赋值 * - test-data - TEST_DATA_PATH - string - 测试集路径 - 平台赋值 * - train-label - TRAIN_LABEL_PATH - string - 训练集标签路径 - 平台赋值 * - val-label - VAL_LABEL_PATH - string - 验证集标签路径 - 平台赋值 * - test-label - TEST_LABEL_PATH - string - 测试集标签路径 - 平台赋值 * - label-class - LABEL_CLASS - string - 标签类别映射字典 - 平台赋值 * - label-class-count - LABEL_CLASS_COUNT - int - 标签类别数量 - 平台赋值 * - epochs - - - int - 单次训练迭代次数 - 训练参数配置/任务提交界面传参 * - weights - - - string - 预训练权重路径 - 训练参数配置/任务提交界面传参 * - batch-size - - - int - 单次训练所抓取的数据样本数量 - 训练参数配置/任务提交界面传参 * - learning-rate - - - float - 学习率 - 训练参数配置/任务提交界面传参 * - device - - - string - 模型运行的硬件设备：CPU或GPU - 训练参数配置/任务提交界面传参 使用全局变量后，平台替换示例： :: # 原始训练脚本中的训练方法 def train(args): """ 训练逻辑 """ data_path = "DATA_PATH" workspace_path = "WORK_DIR_PATH" # 实际训练时脚本中的训练方法 def train(args): """ 训练逻辑 """ data_path = "/root/dataset/" workspace_path = "/root/task/" 2.4 预测脚本 ~~~~~~~~~~~~ - 预测脚本模板 :: import argparse import os import sys from pathlib import Path FILE = Path(__file__).resolve() ROOT = FILE.parents[0] # YOLOv5 root directory if str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # add ROOT to PATH ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative def predict(args): """ 预测逻辑 """ pass def parse_args(): """ 预测脚本参数 """ parser = argparse.ArgumentParser() parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path or triton URL") parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)") parser.add_argument("--confidence", type=float, default=0.25, help="confidence threshold") parser.add_argument("--out-dir", default=ROOT / "runs/detect", help="save results to project/name") parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") args = parser.parse_args() return args if __name__ == "__main__": args = parse_args() predict(args) - 参数传递/全局变量替换 平台提供两种方式，对自定义模型预测脚本 predict.py 中所需的参数进行赋值。 方法一：通过args参数赋值： 平台通过传参的方式，给自定义模型预测脚本提供参数。因此，预测脚本 predict.py 需接收平台提交的参数。 方法二：通过全局变量赋值： 预测脚本 predict.py 中，若使用了平台支持的全局变量，在实际预测时，平台会自动将脚本中的全局变量替换为实际的值。 args参数及全局变量示例，如下表所示。 .. list-table:: :widths: 20 20 15 40 20 :header-rows: 1 * - args参数名称 - 全局变量名称 - 参数类型 - 参数说明 - 赋值方式 * - source - SOURCE_PATH - string - 待预测图片路径 - 预测界面传参 * - out-dir - OUT_DIR_PATH - string - 预测结果保存路径 - 平台赋值 * - weights - - - string - 权重文件路径 - 预测界面传参 * - confidence - - - float - 置信度 - 预测界面传参 * - device - - - string - 模型运行的硬件设备：CPU或GPU - 平台赋值 使用全局变量后，平台替换示例： :: # 原始预测脚本中的预测方法 def predict(args): """ 预测逻辑 """ image_path = "SOURCE_PATH" output_path = "OUT_DIR_PATH" # 实际预测时脚本中的预测方法 def predict(args): """ 预测逻辑 """ image_path = "/root/images/1.jpg" output_path = "/root/predict/result" 注： - 预测脚本输出必须为图片格式。即将 bbox 等信息绘制到图片中，生成预测后的图片。 - 预测图片的名称，与待预测的图片名称保持一致。 2.5 训练日志输出 ~~~~~~~~~~~~~~~~ 自定义模型训练过程需输出日志，且训练指标/训练评估指标的结果，需要以 key:value 的格式输出，示例： :: bbox_mAP: 0.0830, bbox_mAP_50: 0.2360, loss: 0.2754, loss_cls: 0.1174, loss_rpn_cls: 0.0025, loss_rpn_bbox: 0.0039 注：其他格式的输出，平台无法获取到任务相关指标。 3、任务提交 ----------- 平台提供两种方式创建自定义任务： - 图像数据—数据集详情页—标签列表—选择开始训练：自定义算法 - 图像数据—创建任务—选择下一步：自定义算法 两种方式都会进入自定义任务参数配置页，填写参数后（可提前在自定义模型列表页，录入训练参数配置模版，此处可选择模板自动回填模板中的参数），点击开始训练按钮，即完成自定义任务的创建。 参数配置及说明如下表所示。 .. list-table:: :widths: 15 10 10 50 :header-rows: 1 * - 参数名称 - 类型 - 是否必填 - 参数说明 * - 任务名称 - 输入框 - 是 - 填写任务名称 * - 算法名称 - 选择框 - 是 - 从已部署成功的自定义模型中筛选 * - 每次训练Epoch - 输入框 - 是 - 单次训练迭代次数。对应 train.py 脚本中的 epochs 参数 * - 训练次数 - 输入框 - 是 - 训练次数 * - 预计训练时间 - 输入框 - 是 - 任务预计训练时长。任务执行时间超出时长，将会被停止（状态：执行超时）。 * - 希望可视化哪些训练指标 - 输入框 - 是 - 模型训练中的训练指标，包括训练评估指标、损失参数等，如 mAP50、loss。 注：此处填写的指标名称，必须要与 train.py 脚本中使用的名称一致。 * - 训练评估指标 - 选择框 - 是 - 从填写的训练指标中，选择一项作为训练评估指标 * - 输入参数 - 输入框 - 是 - 模型训练中所需要的参数。录入：参数名称、参数的值、是否进行搜索。 注：参数的值，可为单个值，或范围。示例：10 或 [5,10]。若选择了进行搜索，则需要填写范围。示例：[5，10]。 4、权重下载 ----------- 自定义任务状态包括：执行中、执行完成、执行失败、用户停止、执行超时、等待执行。 执行完成的自定义任务支持下载权重。 5、模型预测 ----------- 模型预测页面，选择训练好的模型，上传待预测的图片，即能得到预测结果。