自定义模型规则#

1、数据准备#

平台支持8种任务类型的自定义模型训练，包括目标检测、语义分割、实例分割、图像分类、文字检测、文字识别、旋转目标检测、其他。

自定义模型使用平台提供的数据集进行训练。不同任务类型的数据集，拆分格式不同，详情请见图像数据拆分格式。

用户选择任务类型：其他，创建新的数据集，平台只解析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 写法：

直接将相关依赖项写入 requirements.txt 文件中，文件内容如下所示：

matplotlib>=3.3
numpy>=1.23.5
opencv-python>=4.1.1
pillow>=10.3.0
...

模型安装所需依赖项写在其他文件中，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参数及全局变量示例，如下表所示。

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参数及全局变量示例，如下表所示。

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、任务提交#

平台提供两种方式创建自定义任务：

图像数据—数据集详情页—标签列表—选择开始训练：自定义算法
图像数据—创建任务—选择下一步：自定义算法

两种方式都会进入自定义任务参数配置页，填写参数后（可提前在自定义模型列表页，录入训练参数配置模版，此处可选择模板自动回填模板中的参数），点击开始训练按钮，即完成自定义任务的创建。

参数配置及说明如下表所示。

参数名称	类型	是否必填	参数说明
任务名称	输入框	是	填写任务名称
算法名称	选择框	是	从已部署成功的自定义模型中筛选
每次训练Epoch	输入框	是	单次训练迭代次数。对应 train.py 脚本中的 epochs 参数
训练次数	输入框	是	训练次数
预计训练时间	输入框	是	任务预计训练时长。任务执行时间超出时长，将会被停止（状态：执行超时）。
希望可视化哪些训练指标	输入框	是	模型训练中的训练指标，包括训练评估指标、损失参数等，如 mAP50、loss。注：此处填写的指标名称，必须要与 train.py 脚本中使用的名称一致。
训练评估指标	选择框	是	从填写的训练指标中，选择一项作为训练评估指标
输入参数	输入框	是	模型训练中所需要的参数。录入：参数名称、参数的值、是否进行搜索。注：参数的值，可为单个值，或范围。示例：10 或 [5,10]。若选择了进行搜索，则需要填写范围。示例：[5，10]。

4、权重下载#

自定义任务状态包括：执行中、执行完成、执行失败、用户停止、执行超时、等待执行。

执行完成的自定义任务支持下载权重。

5、模型预测#

模型预测页面，选择训练好的模型，上传待预测的图片，即能得到预测结果。

Coovally 用户手册 v2.1.2

自定义模型规则

Contents