6.1 KiB
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.yaml或pipeline.config。
3.4 模型训练 (Training)
- 异步训练: 后端通过
subprocess启动独立的训练进程,避免阻塞主线程。 - 实时日志: 利用 Django Channels + WebSocket 实现训练日志实时推送到前端展示。
- 参数配置: 支持自定义 Epochs, Batch Size, Learning Rate 等超参数。
- 断点续训与终止: 支持训练任务的中断与状态回滚。
3.5 在线预测 (Prediction)
- 即时推理: 上传图片后,后端加载对应模型进行推理并返回标注后的图片或坐标数据。
- 多模型支持: 根据
algorithm_engine自动选择 YOLO 或 D-FINE 推理器。
4. 关键代码逻辑说明
4.1 训练流程
- 创建记录:
create_algorithm_train_record初始化训练任务。 - 数据准备:
prepare_algorithm_train_dataset将数据集样本复制到训练目录,并生成 YOLO/COCO 格式标签。 - 启动训练:
start_algorithm_train根据引擎类型 (train_algorithm或train_algorithm_dfine) 启动子进程。 - 日志流:
stream_logs_to_frontend线程实时读取日志文件并通过 WebSocket 发送。 - 模型归档: 训练完成后,
archive_train_model将最佳权重文件 (best.pt) 归档为新模型版本。
4.2 遗留与待优化项 (Refactoring Notes)
- prediction_server_legacy: 该应用名包含 "legacy",但实际上承载了所有核心 Models (
AiAlgorithm,AiAlgorithmTrainRecords等)。建议在未来重构中重命名为core或base。 - 冗余代码:
backend/apps/algorithms/algorithm/下存在多个未使用的目录(如tensorflow,other,dataPrediction),建议在确认无业务依赖后移除。 - 配置管理: 根目录下的
settings.py为废弃文件,实际配置位于backend/config/settings/。
5. 部署与环境
5.1 依赖安装
cd backend
pip install -r ../requirements/all.txt
注意:需确保安装了正确版本的 torch 和 ultralytics 以支持 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 和现代化前端技术,提升了用户体验(如实时日志)。后续建议清理遗留代码并优化应用命名,以提升代码可维护性。