model_train_dm/docs/系统技术方案说明.md
2026-07-27 17:51:49 +08:00

6.1 KiB

AI 算法模型训练及部署管理平台 - 系统技术方案说明

1. 项目概述

本项目是一个集成的 AI 算法模型全生命周期管理平台,旨在提供从数据管理、标注、模型训练到在线预测的一站式解决方案。项目基于前后端分离架构,前端采用现代化的 Vue 3 技术栈,后端基于 Django 框架构建,支持 YOLO 和 D-FINE 等主流目标检测算法的训练与推理。

2. 系统架构

2.1 技术栈

前端 (Frontend)

  • 框架: Vue 3 (Composition API)
  • 构建工具: Vite 4
  • UI 组件库: Element Plus
  • 状态管理: Pinia
  • 样式: TailwindCSS + SCSS
  • 路由: Vue Router 4
  • HTTP 客户端: Axios
  • 可视化: ECharts
  • 标注工具: ailabel

后端 (Backend)

  • 框架: Django 4.2.7 + Django REST Framework (DRF)
  • 实时通信: Django Channels (WebSocket)
  • 数据库: PostgreSQL (pgsqldb)
  • 深度学习框架: PyTorch (主要), Ultralytics (YOLO), PaddlePaddle (遗留/特定模块)
  • 加密: SM4 (国密算法)
  • 任务调度: Threading / Subprocess (当前实现), 可扩展 Celery

运行环境

  • 操作系统: Windows (开发环境) / Linux (生产环境)
  • Python 版本: 3.11
  • 依赖管理: pip + virtualenv

2.2 目录结构

AI_TrainPrediction/
├── backend/                  # 后端项目根目录
│   ├── apps/                 # 业务应用模块
│   │   ├── algorithms/       # 算法管理 (模型定义, 训练逻辑, 预测接口)
│   │   ├── annotation/       # 标注管理
│   │   ├── common/           # 公共工具 (序列化, 权限, 文件处理)
│   │   ├── datasets/         # 数据集管理
│   │   ├── iam/              # 身份认证与权限 (用户, 组织, 角色)
│   │   ├── core/             # 核心数据模型 (数据集, 样本, 模型, 训练任务)
│   │   └── training/         # 训练任务管理 (日志流, 进程控制)
│   ├── config/               # Django 配置 (settings, urls, asgi/wsgi)
│   ├── db/                   # 数据库相关
│   ├── enev/                 # Python 虚拟环境
│   └── manage.py             # Django 管理脚本
├── frontend/                 # 前端项目根目录
│   ├── src/                  # 源代码
│   ├── public/               # 静态资源
│   └── vite.config.ts        # Vite 配置
├── requirements/             # Python 依赖清单
└── tools/                    # 辅助脚本

3. 核心功能模块

3.1 身份认证与权限 (IAM)

  • JWT 认证: 使用 simplejwt 实现基于 Token 的身份验证。
  • SM4 加密: 登录密码传输采用国密 SM4 算法加密,保障安全性。
  • RBAC: 基于 Django 内置的 Group 和 Permission 实现角色权限控制。

3.2 数据集管理 (Datasets)

  • 多格式支持: 支持图片上传、解压、预览。
  • 标注集成: 集成 ailabel 前端库,支持在线标注并生成 XML/JSON 格式标注文件。
  • 样本管理: AiDatasetSample 表记录样本状态、路径及标注内容。

3.3 算法与模型管理 (Algorithms)

  • 算法引擎:
    • 01 (YOLO): 基于 Ultralytics 框架,支持 YOLOv8 等模型。
    • 02 (D-FINE): 基于 PyTorch 实现的改进版 RT-DETR。
  • 模型版本控制: 自动生成版本号,支持模型归档、导出。
  • 配置文件生成: 根据训练参数动态生成 config.yamlpipeline.config

3.4 模型训练 (Training)

  • 异步训练: 后端通过 subprocess 启动独立的训练进程,避免阻塞主线程。
  • 实时日志: 利用 Django Channels + WebSocket 实现训练日志实时推送到前端展示。
  • 参数配置: 支持自定义 Epochs, Batch Size, Learning Rate 等超参数。
  • 断点续训与终止: 支持训练任务的中断与状态回滚。

3.5 在线预测 (Prediction)

  • 即时推理: 上传图片后,后端加载对应模型进行推理并返回标注后的图片或坐标数据。
  • 多模型支持: 根据 algorithm_engine 自动选择 YOLO 或 D-FINE 推理器。

4. 关键代码逻辑说明

4.1 训练流程

  1. 创建记录: create_algorithm_train_record 初始化训练任务。
  2. 数据准备: prepare_algorithm_train_dataset 将数据集样本复制到训练目录,并生成 YOLO/COCO 格式标签。
  3. 启动训练: start_algorithm_train 根据引擎类型 (train_algorithmtrain_algorithm_dfine) 启动子进程。
  4. 日志流: stream_logs_to_frontend 线程实时读取日志文件并通过 WebSocket 发送。
  5. 模型归档: 训练完成后,archive_train_model 将最佳权重文件 (best.pt) 归档为新模型版本。

4.2 遗留与待优化项 (Refactoring Notes)

  • prediction_server_legacy: 该应用名包含 "legacy",但实际上承载了所有核心 Models (AiAlgorithm, AiAlgorithmTrainRecords 等)。建议在未来重构中重命名为 corebase
  • 冗余代码: backend/apps/algorithms/algorithm/ 下存在多个未使用的目录(如 tensorflow, other, dataPrediction),建议在确认无业务依赖后移除。
  • 配置管理: 根目录下的 settings.py 为废弃文件,实际配置位于 backend/config/settings/

5. 部署与环境

5.1 依赖安装

cd backend
pip install -r ../requirements/all.txt

注意:需确保安装了正确版本的 torchultralytics 以支持 GPU 训练。

5.2 数据库配置

修改 backend/config/settings/base.py 中的 DATABASES 配置以连接 PostgreSQL。

5.3 启动服务

# 后端
python manage.py runserver 0.0.0.0:8000

# 前端
cd frontend
npm install
npm run dev

6. 总结

本项目已完成从旧系统的核心功能迁移,后端服务稳定,能够支撑完整的 AI 模型训练与应用流程。通过引入 Django Channels 和现代化前端技术,提升了用户体验(如实时日志)。后续建议清理遗留代码并优化应用命名,以提升代码可维护性。