从勘智云平台拆解：如何自建一套端到端AI模型训练与部署平台

1
2
3
4

创建数据集 → 上传/标注数据 → 配置训练参数 → 下载部署包
     │              │               │              │
     ▼              ▼               ▼              ▼
  选择任务类型    在线标注工具     GPU集群训练    KModel+部署脚本

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# 克隆CVAT仓库
git clone https://github.com/cvat-ai/cvat.git
cd cvat

# 使用docker-compose启动
docker compose up -d

# 创建超级用户（进入容器执行）
docker exec -it cvat_server bash -c \
  'python3 ~/manage.py createsuperuser'

# 访问 http://localhost:8080

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

import redis
import json

class TrainingQueue:
    """训练任务队列，基于Redis实现"""
    
    def __init__(self, redis_url="redis://localhost:6379"):
        self.r = redis.from_url(redis_url)
        self.queue_key = "training:queue"
        self.running_key = "training:running"
    
    def submit(self, task_id: str, config: dict, user_id: str):
        """提交训练任务到队列
        
        Args:
            task_id: 任务唯一ID
            config: 训练配置（模型类型、超参数等）
            user_id: 提交用户
        """
        # 检查用户是否已有任务在运行
        running_tasks = self.r.hgetall(self.running_key)
        for tid, data in running_tasks.items():
            task = json.loads(data)
            if task["user_id"] == user_id:
                return {"error": "每个用户同时只能运行1个任务"}
        
        # 入队
        self.r.rpush(self.queue_key, json.dumps({
            "task_id": task_id,
            "config": config,
            "user_id": user_id,
            "status": "queued"
        }))
        return {"task_id": task_id, "status": "queued"}
    
    def fetch_next(self, gpu_id: str) -> dict | None:
        """GPU空闲时拉取下一个任务"""
        task_data = self.r.lpop(self.queue_key)
        if task_data is None:
            return None
        task = json.loads(task_data)
        task["status"] = "running"
        task["gpu_id"] = gpu_id
        self.r.hset(self.running_key, task["task_id"], 
                     json.dumps(task))
        return task
    
    def complete(self, task_id: str, result: dict):
        """任务完成，从运行列表移除"""
        self.r.hdel(self.running_key, task_id)
        self.r.set(
            f"training:result:{task_id}",
            json.dumps(result),
            ex=86400  # 结果保留24小时
        )

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import models, transforms
from torchvision.datasets import ImageFolder

def train_classification(data_dir: str, 
                          epochs: int = 50,
                          lr: float = 0.001,
                          num_classes: int = 10):
    """图像分类训练函数
    
    对应勘智云"图像分类"任务：
    - 数据格式：data_dir下按类别文件夹组织
    - 模型：ResNet18，适合K210/K230算力
    - 输出：PyTorch .pt模型文件
    
    Args:
        data_dir: 数据集目录，子文件夹名即类别名
        epochs: 训练轮数
        lr: 学习率
        num_classes: 类别数
    """
    # 数据预处理
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])
    
    # 加载数据集（ImageFolder自动按文件夹名分类）
    train_set = ImageFolder(
        f"{data_dir}/train", transform=transform
    )
    val_set = ImageFolder(
        f"{data_dir}/val", transform=transform
    )
    train_loader = DataLoader(
        train_set, batch_size=32, shuffle=True
    )
    val_loader = DataLoader(
        val_set, batch_size=32, shuffle=False
    )
    
    # 模型选择：ResNet18（轻量，适合边缘设备）
    model = models.resnet18(
        weights=models.ResNet18_Weights.DEFAULT
    )
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    model = model.cuda()
    
    # 训练配置
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(
        model.parameters(), lr=lr
    )
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer, step_size=10, gamma=0.5
    )
    
    # 训练循环
    best_acc = 0.0
    for epoch in range(epochs):
        # 训练阶段
        model.train()
        for images, labels in train_loader:
            images, labels = images.cuda(), labels.cuda()
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        # 验证阶段
        model.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.cuda(), labels.cuda()
                outputs = model(images)
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        acc = correct / total
        if acc > best_acc:
            best_acc = acc
            torch.save(
                model.state_dict(),
                "best_model.pt"
            )
        
        scheduler.step()
        print(f"Epoch {epoch+1}/{epochs} "
              f"Val Acc: {acc:.4f} Best: {best_acc:.4f}")
    
    return best_acc

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

# 训练脚本中，每个epoch结束后推送loss
import websockets
import asyncio
import json

async def push_metrics(task_id: str, epoch: int, 
                        loss: float, acc: float):
    """通过WebSocket推送训练指标到前端"""
    async with websockets.connect(
        f"ws://localhost:8765/ws/{task_id}"
    ) as ws:
        await ws.send(json.dumps({
            "type": "metrics",
            "epoch": epoch,
            "loss": loss,
            "acc": acc,
            "timestamp": int(time.time())
        }))

# 在训练循环中调用
# await push_metrics(task_id, epoch, loss_val, acc_val)

1
2
3
4
5

训练好的模型          量化过程               部署模型
┌──────────┐      ┌──────────────┐      ┌──────────┐
│ float32  │ ───→ │ 量化到int8   │ ───→ │ .kmodel  │
│ ~45MB    │      │ 精度损失<2%  │      │ ~11MB    │
└──────────┘      └──────────────┘      └──────────┘

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

# 完整的模型转换链路：PyTorch → ONNX → TFLite → 边缘设备

import torch

# 第一步：PyTorch模型导出为ONNX
def export_to_onnx(model, input_shape=(1, 3, 224, 224),
                    output_path="model.onnx"):
    """将PyTorch模型导出为ONNX格式
    
    ONNX是模型转换的"中间语言"，
    几乎所有推理框架都能读ONNX。
    """
    model.eval()
    dummy = torch.randn(*input_shape)
    torch.onnx.export(
        model, dummy, output_path,
        export_params=True,
        opset_version=14,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {0: "batch"},
            "output": {0: "batch"}
        }
    )
    print(f"ONNX模型已保存: {output_path}")

# 第二步：ONNX转换为TFLite（量化）
def onnx_to_tflite_int8(onnx_path="model.onnx",
                          output_path="model_int8.tflite",
                          rep_data_dir="data/train"):
    """ONNX转TFLite并执行int8量化
    
    需要提供代表性数据集（rep_data）来校准量化参数。
    这和勘智云的KModel转换原理一致：
    都是用校准数据确定每层的量化scale和zero_point。
    """
    import onnx
    from onnx_tf.backend import prepare
    import tensorflow as tf
    
    # ONNX → TF
    onnx_model = onnx.load(onnx_path)
    tf_rep = prepare(onnx_model)
    tf_rep.export_graph("model_tf")
    
    # 收集校准数据
    rep_data = []
    for img_path in glob.glob(
        f"{rep_data_dir}/**/*.jpg", recursive=True
    )[:100]:  # 100张足够校准
        img = preprocess_image(img_path)  # resize+normalize
        rep_data.append(img)
    
    # TF → TFLite (int8量化)
    converter = tf.lite.TFLiteConverter.from_saved_model("model_tf")
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = lambda: rep_data
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS_INT8
    ]
    converter.inference_input_type = tf.uint8
    converter.inference_output_type = tf.uint8
    
    tflite_model = converter.convert()
    with open(output_path, "wb") as f:
        f.write(tflite_model)
    print(f"量化TFLite模型已保存: {output_path}")
    print(f"模型大小: {len(tflite_model)/1024/1024:.1f}MB")

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

import zipfile
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.base import MIMEBase
from email import encoders
from pathlib import Path
import json

def generate_deployment_package(
    task_id: str,
    model_path: str,
    task_type: str,
    labels: list,
    output_dir: str = "/tmp/deploy"
) -> str:
    """生成部署包（ZIP格式）
    
    模拟勘智云的部署包结构：
    ├── deploy_config.json
    ├── model.kmodel
    ├── mp_deployment_source/
    │   ├── infer.py
    │   └── __init__.py
    └── README.txt
    """
    deploy_dir = Path(output_dir) / task_id
    deploy_dir.mkdir(parents=True, exist_ok=True)
    
    # 1. 拷贝模型文件
    import shutil
    shutil.copy(model_path, deploy_dir / "model.kmodel")
    
    # 2. 生成部署配置
    deploy_config = {
        "task_type": task_type,
        "model_file": "model.kmodel",
        "labels": labels,
        "thresholds": {
            "confidence": 0.5,
            "nms": 0.45
        },
        "input_size": [224, 224],
        "mean": [0.485, 0.456, 0.406],
        "std": [0.229, 0.224, 0.225]
    }
    with open(deploy_dir / "deploy_config.json", "w") as f:
        json.dump(deploy_config, f, indent=2, ensure_ascii=False)
    
    # 3. 生成MicroPython推理脚本
    infer_code = '''# MicroPython推理脚本
import json
import sensor
import image
import KPU  # K230内置的推理库

# 加载配置
with open("deploy_config.json") as f:
    config = json.load(f)

# 加载模型
task = kpu.load(config["model_file"])
kpu.init_yolo2(task, config["thresholds"]["confidence"], 
               config["thresholds"]["nms"])

# 摄像头初始化
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.run(1)

while True:
    img = sensor.snapshot()
    # 预处理
    img = img.resize(224, 224)
    # 推理
    results = kpu.yolo2_run(task, img)
    for obj in results:
        # 在画面上绘制结果
        img.draw_rectangle(obj.x(), obj.y(), 
                          obj.w(), obj.h())
    print(results)
'''
    mp_dir = deploy_dir / "mp_deployment_source"
    mp_dir.mkdir(exist_ok=True)
    (mp_dir / "infer.py").write_text(infer_code)
    
    # 4. 打包ZIP
    zip_path = f"{output_dir}/{task_id}_deploy.zip"
    with zipfile.ZipFile(zip_path, "w", 
                          zipfile.ZIP_DEFLATED) as zf:
        for file in deploy_dir.rglob("*"):
            if file.is_file():
                zf.write(file, 
                         file.relative_to(deploy_dir))
    
    return zip_path


def send_deployment_email(
    to_email: str, 
    task_id: str,
    zip_path: str
):
    """训练完成后发送部署包到用户邮箱"""
    msg = MIMEMultipart()
    msg["Subject"] = f"训练任务 {task_id} 已完成"
    msg["From"] = "[email protected]"
    msg["To"] = to_email
    
    body = f"""
    您的训练任务 {task_id} 已完成！
    
    部署包已作为附件发送，解压后请阅读README.txt进行部署。
    
    部署步骤：
    1. 将 model.kmodel 复制到开发板
    2. 将 deploy_config.json 复制到开发板
    3. 运行 mp_deployment_source/infer.py
    4. 如效果不佳，调整 deploy_config.json 中的阈值
    """
    msg.attach(MIMEText(body, "plain", "utf-8"))
    
    with open(zip_path, "rb") as f:
        part = MIMEBase("application", "zip")
        part.set_payload(f.read())
        encoders.encode_base64(part)
        part.add_header(
            "Content-Disposition",
            f"attachment; filename={task_id}_deploy.zip"
        )
        msg.attach(part)
    
    # 发送邮件（实际环境用SMTP配置）
    with smtplib.SMTP("smtp.example.com", 587) as server:
        server.starttls()
        server.login("user", "password")
        server.send_message(msg)
    
    print(f"部署包已发送至 {to_email}")

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110

# api/dataset.py — 数据集上传接口
from fastapi import APIRouter, UploadFile, File, BackgroundTasks
from pathlib import Path
import zipfile, shutil, json

router = APIRouter(prefix="/api/dataset", tags=["dataset"])

@router.post("/upload")
async def upload_dataset(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    task_type: str = "classification",  # classification|detection|segmentation
    user_id: str = "user_001"
):
    """上传数据集压缩包
    
    对应勘智云的"上传训练数据"步骤：
    1. 接收zip文件
    2. 后台解压并校验
    3. 注册到数据集表
    4. 返回dataset_id供后续训练使用
    """
    # 保存上传的zip文件
    upload_dir = Path(f"/data/uploads/{user_id}")
    upload_dir.mkdir(parents=True, exist_ok=True)
    zip_path = upload_dir / file.filename
    with open(zip_path, "wb") as f:
        shutil.copyfileobj(file.file, f)
    
    # 后台异步解压+校验（不阻塞响应）
    dataset_id = f"ds_{int(time.time())}"
    background_tasks.add_task(
        extract_and_validate, zip_path, dataset_id, task_type
    )
    
    return {"dataset_id": dataset_id, "status": "extracting"}


def extract_and_validate(zip_path: Path, dataset_id: str, 
                          task_type: str):
    """解压数据集并校验格式（后台任务）
    
    不同任务类型的校验规则：
    - classification: 检查子文件夹名是否为类别名，每类至少10张图
    - detection: 检查每张图片是否有同名XML标注文件
    - segmentation: 检查每张图片是否有同名JSON标注文件
    """
    dataset_dir = Path(f"/data/datasets/{dataset_id}")
    dataset_dir.mkdir(parents=True, exist_ok=True)
    
    # 解压
    with zipfile.ZipFile(zip_path, "r") as zf:
        zf.extractall(dataset_dir)
    
    # 校验
    image_exts = {".jpg", ".jpeg", ".png", ".bmp"}
    images = list(dataset_dir.rglob("*"))
    images = [f for f in images if f.suffix.lower() in image_exts]
    
    if len(images) == 0:
        update_dataset_status(dataset_id, "error", 
                              "未找到图片文件")
        return
    if len(images) > 1000:
        update_dataset_status(dataset_id, "error",
                              f"图片数量{len(images)}超过1000上限")
        return
    
    # 按任务类型校验标注文件
    if task_type == "detection":
        for img in images:
            xml_file = img.with_suffix(".xml")
            if not xml_file.exists():
                update_dataset_status(dataset_id, "error",
                    f"检测标注缺失: {img.name} 需要同名XML文件")
                return
    elif task_type == "segmentation":
        for img in images:
            json_file = img.with_suffix(".json")
            if not json_file.exists():
                update_dataset_status(dataset_id, "error",
                    f"分割标注缺失: {img.name} 需要同名JSON文件")
                return
    
    # 校验通过，统计类别
    if task_type == "classification":
        # 子文件夹名就是类别名
        classes = set(
            f.parent.name for f in images
            if f.parent != dataset_dir
        )
        class_list = sorted(classes)
    else:
        class_list = ["object"]  # 检测/分割只有一类目标
    
    # 写入数据集元信息
    meta = {
        "dataset_id": dataset_id,
        "task_type": task_type,
        "num_images": len(images),
        "classes": class_list,
        "num_classes": len(class_list),
        "path": str(dataset_dir),
        "status": "ready"
    }
    with open(f"/data/datasets/{dataset_id}/meta.json", "w") as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)
    
    update_dataset_status(dataset_id, "ready", 
                         f"{len(images)}张图片, {len(class_list)}个类别")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

# api/training.py — 训练任务提交接口
from fastapi import APIRouter
from pydantic import BaseModel

router = APIRouter(prefix="/api/training", tags=["training"])

class TrainingRequest(BaseModel):
    """训练请求参数 — 对应勘智云的训练配置页面"""
    dataset_id: str           # 上一步拿到的数据集ID
    task_type: str            # classification|detection|segmentation
    epochs: int = 50          # 训练轮数
    learning_rate: float = 0.001  # 学习率
    batch_size: int = 32      # 批大小
    backbone: str = "resnet18"  # 骨干网络
    augmentation: bool = True   # 是否开启数据增强
    user_email: str           # 训练完成后通知的邮箱

@router.post("/submit")
async def submit_training(req: TrainingRequest):
    """提交训练任务
    
    流程：校验数据集 → 检查用户配额 → 入队列 → 返回task_id
    """
    # 1. 读取数据集元信息
    meta = load_dataset_meta(req.dataset_id)
    if meta["status"] != "ready":
        return {"error": f"数据集未就绪: {meta['status']}"}
    
    # 2. 检查用户是否已有任务在运行（勘智云的限制：每用户1个并发）
    queue = TrainingQueue()
    check = queue.submit(
        task_id=f"task_{int(time.time())}",
        config={
            "dataset_id": req.dataset_id,
            "task_type": req.task_type,
            "epochs": req.epochs,
            "learning_rate": req.learning_rate,
            "batch_size": req.batch_size,
            "backbone": req.backbone,
            "augmentation": req.augmentation,
            "num_classes": meta["num_classes"],
            "data_path": meta["path"],
            "classes": meta["classes"],
            "user_email": req.user_email
        },
        user_id="user_001"  # 实际从JWT token解析
    )
    
    if "error" in check:
        return check
    
    return {"task_id": check["task_id"], "status": "queued"}

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313

# worker/training_worker.py — GPU训练工作进程
import time, json, subprocess, os

def gpu_worker_loop(gpu_id: int = 0):
    """GPU Worker主循环
    
    这个进程常驻运行，不停从Redis队列拉任务。
    勘智云有8块V100，意味着可以同时跑8个worker。
    每个worker绑定一块GPU。
    """
    queue = TrainingQueue()
    
    while True:
        # 尝试拉取任务
        task = queue.fetch_next(gpu_id=str(gpu_id))
        if task is None:
            time.sleep(5)  # 没有任务，等5秒再查
            continue
        
        # 有任务了，开始训练
        print(f"[GPU {gpu_id}] 拉取任务 {task['task_id']}, "
              f"类型: {task['config']['task_type']}")
        
        run_training_pipeline(task, gpu_id)


def run_training_pipeline(task: dict, gpu_id: int):
    """完整的训练流水线 — 串起所有环节
    
    这就是勘智云"提交训练"按钮背后执行的完整流程：
    1. 启动WebSocket日志服务
    2. 运行训练脚本（根据task_type选择）
    3. 训练完成后导出ONNX
    4. 执行int8量化
    5. 生成部署包
    6. 发送邮件通知
    """
    task_id = task["task_id"]
    config = task["config"]
    
    # ---------- 阶段1：数据准备 ----------
    update_task_status(task_id, "preparing_data", 
                       "正在准备训练数据...")
    
    # 如果开启数据增强，生成增强后的数据集
    if config.get("augmentation", False):
        augmented_dir = apply_data_augmentation(
            config["data_path"],
            task_type=config["task_type"]
        )
        train_data_dir = augmented_dir
    else:
        # 自动划分训练集/验证集 (8:2)
        train_data_dir = split_train_val(
            config["data_path"], ratio=0.8
        )
    
    # ---------- 阶段2：启动训练 ----------
    update_task_status(task_id, "training",
                       f"训练中... {config['epochs']}轮")
    
    # 根据任务类型选择训练脚本
    # 训练脚本内部就是前面"模块二"里的train_classification等函数
    task_type = config["task_type"]
    
    if task_type == "classification":
        from training_scripts import train_classification
        best_acc, model_path = train_classification(
            data_dir=train_data_dir,
            epochs=config["epochs"],
            lr=config["learning_rate"],
            num_classes=config["num_classes"],
            backbone=config["backbone"],
            task_id=task_id  # 传入task_id用于推送日志
        )
    elif task_type == "detection":
        from training_scripts import train_detection
        best_map, model_path = train_detection(
            data_dir=train_data_dir,
            epochs=config["epochs"],
            lr=config["learning_rate"],
            num_classes=config["num_classes"],
            backbone=config["backbone"],
            task_id=task_id
        )
    elif task_type == "segmentation":
        from training_scripts import train_segmentation
        best_miou, model_path = train_segmentation(
            data_dir=train_data_dir,
            epochs=config["epochs"],
            lr=config["learning_rate"],
            num_classes=config["num_classes"],
            backbone=config["backbone"],
            task_id=task_id
        )
    
    # ---------- 阶段3：模型评估 ----------
    update_task_status(task_id, "evaluating",
                       "正在生成评估报告...")
    
    eval_report = evaluate_model(
        model_path=model_path,
        data_dir=config["data_path"],
        task_type=task_type,
        classes=config["classes"]
    )
    # eval_report包含：准确率/召回率/F1、混淆矩阵、PR曲线数据
    
    save_eval_report(task_id, eval_report)
    
    # ---------- 阶段4：模型转换与量化 ----------
    update_task_status(task_id, "converting",
                       "正在导出ONNX并量化...")
    
    # PyTorch → ONNX
    onnx_path = export_to_onnx(
        model=model_path,
        output_path=f"/data/output/{task_id}/model.onnx"
    )
    
    # ONNX → TFLite int8
    tflite_path = onnx_to_tflite_int8(
        onnx_path=onnx_path,
        output_path=f"/data/output/{task_id}/model_int8.tflite",
        rep_data_dir=config["data_path"]
    )
    
    # ---------- 阶段5：生成部署包 ----------
    update_task_status(task_id, "packaging",
                       "正在生成部署包...")
    
    zip_path = generate_deployment_package(
        task_id=task_id,
        model_path=tflite_path,
        task_type=task_type,
        labels=config["classes"]
    )
    
    # ---------- 阶段6：通知用户 ----------
    send_deployment_email(
        to_email=config["user_email"],
        task_id=task_id,
        zip_path=zip_path
    )
    
    # 标记任务完成
    queue.complete(task_id, {
        "status": "completed",
        "best_metric": eval_report.get("best_metric"),
        "model_path": zip_path,
        "eval_report_id": task_id
    })
    
    update_task_status(task_id, "completed",
                       f"训练完成！准确率: {eval_report.get('best_metric', 0):.2%}")
    print(f"[GPU {gpu_id}] 任务 {task_id} 完成")


def apply_data_augmentation(data_path: str, 
                             task_type: str) -> str:
    """数据增强 — 在1000张图片限制下提升效果
    
    勘智云限制1000张图片，数据增强就特别重要：
    原始1000张 → 增强5000张，相当于5倍数据量。
    
    策略：
    - 水平翻转（最安全，不改变语义）
    - 随机裁剪（模拟不同拍摄角度）
    - 色彩抖动（模拟不同光照条件）
    - 不用垂直翻转（鸟倒过来飞不合理）
    """
    from torchvision import transforms
    from PIL import Image
    import shutil
    
    aug_dir = f"{data_path}_augmented"
    shutil.copytree(data_path, aug_dir, dirs_exist_ok=True)
    
    aug_transform = transforms.Compose([
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
        transforms.ColorJitter(
            brightness=0.2, contrast=0.2,
            saturation=0.2, hue=0.1
        ),
    ])
    
    image_exts = {".jpg", ".jpeg", ".png"}
    for img_file in Path(data_path).rglob("*"):
        if img_file.suffix.lower() in image_exts:
            img = Image.open(img_file).convert("RGB")
            # 每张图生成3个增强版本
            for i in range(3):
                augmented = aug_transform(img)
                aug_path = (Path(aug_dir) / img_file.parent 
                           / f"{img_file.stem}_aug{i}{img_file.suffix}")
                augmented.save(aug_path)
    
    return aug_dir


def split_train_val(data_path: str, ratio: float = 0.8) -> str:
    """自动划分训练集和验证集
    
    把原始数据按8:2分成 train/ 和 val/ 两个子目录。
    训练脚本会自动从这两个目录读取数据。
    """
    import random, shutil
    random.seed(42)  # 固定随机种子，保证可复现
    
    images = list(Path(data_path).rglob("*"))
    images = [f for f in images 
              if f.suffix.lower() in {".jpg", ".jpeg", ".png"}]
    random.shuffle(images)
    
    split_idx = int(len(images) * ratio)
    train_imgs = images[:split_idx]
    val_imgs = images[split_idx:]
    
    for split_name, split_imgs in [("train", train_imgs), 
                                    ("val", val_imgs)]:
        split_dir = Path(data_path) / split_name
        split_dir.mkdir(exist_ok=True)
        for img in split_imgs:
            # 保持子目录结构（分类任务按类别文件夹）
            rel = img.relative_to(Path(data_path))
            dst = split_dir / rel
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(img, dst)
    
    return data_path  # 返回数据根目录，训练函数自己找train/和val/


def evaluate_model(model_path: str, data_dir: str,
                   task_type: str, classes: list) -> dict:
    """模型评估 — 生成完整的评估报告
    
    不同任务类型的评估指标：
    - classification: Top-1准确率、混淆矩阵
    - detection: [email protected]、PR曲线
    - segmentation: mIoU、每类IoU
    
    这里以分类任务为例，生成混淆矩阵。
    """
    import torch
    from torchvision import transforms, models
    from torchvision.datasets import ImageFolder
    from torch.utils.data import DataLoader
    import numpy as np
    from sklearn.metrics import (confusion_matrix, 
                                  classification_report)
    
    # 加载训练好的模型
    model = models.resnet18()
    model.fc = torch.nn.Linear(model.fc.in_features, len(classes))
    model.load_state_dict(torch.load(model_path, 
                                     map_location="cpu"))
    model.eval()
    
    # 准备验证集
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])
    
    val_dir = f"{data_dir}/val"
    if not Path(val_dir).exists():
        val_dir = data_dir  # 没有划分的话用整个数据集
    
    val_set = ImageFolder(val_dir, transform=transform)
    val_loader = DataLoader(val_set, batch_size=32, shuffle=False)
    
    all_preds, all_labels = [], []
    with torch.no_grad():
        for images, labels in val_loader:
            outputs = model(images)
            _, preds = outputs.max(1)
            all_preds.extend(preds.numpy())
            all_labels.extend(labels.numpy())
    
    # 生成混淆矩阵
    cm = confusion_matrix(all_labels, all_preds)
    
    # 分类报告（精确率、召回率、F1）
    report = classification_report(
        all_labels, all_preds,
        target_names=classes,
        output_dict=True  # 返回字典格式
    )
    
    accuracy = report["accuracy"]
    
    return {
        "task_type": task_type,
        "best_metric": accuracy,
        "accuracy": accuracy,
        "confusion_matrix": cm.tolist(),
        "per_class_report": {
            cls: {
                "precision": report[cls]["precision"],
                "recall": report[cls]["recall"],
                "f1_score": report[cls]["f1-score"],
                "support": report[cls]["support"]
            }
            for cls in classes if cls in report
        },
        "macro_avg": report.get("macro avg", {}),
        "weighted_avg": report.get("weighted avg", {})
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

# websocket/training_ws.py — WebSocket服务端
from fastapi import WebSocket, WebSocketDisconnect
import json, asyncio

class ConnectionManager:
    """管理WebSocket连接 — 每个训练任务一个频道"""
    
    def __init__(self):
        # task_id → [WebSocket连接列表]
        self.connections: dict[str, list[WebSocket]] = {}
    
    async def connect(self, websocket: WebSocket, task_id: str):
        await websocket.accept()
        if task_id not in self.connections:
            self.connections[task_id] = []
        self.connections[task_id].append(websocket)
    
    def disconnect(self, websocket: WebSocket, task_id: str):
        if task_id in self.connections:
            self.connections[task_id].remove(websocket)
    
    async def broadcast(self, task_id: str, message: dict):
        """向所有监听该任务的前端推送消息"""
        if task_id not in self.connections:
            return
        dead = []
        for ws in self.connections[task_id]:
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.connections[task_id].remove(ws)

manager = ConnectionManager()

# FastAPI WebSocket路由
@router.websocket("/ws/{task_id}")
async def training_websocket(websocket: WebSocket, 
                              task_id: str):
    """训练实时日志WebSocket
    
    前端连接后能收到：
    - {"type": "status", "status": "training", "msg": "..."}
    - {"type": "metrics", "epoch": 1, "loss": 0.85, "acc": 0.72}
    - {"type": "eval", "accuracy": 0.95, "confusion_matrix": [...]}
    - {"type": "done", "download_url": "/api/download/task_xxx"}
    """
    await manager.connect(websocket, task_id)
    try:
        while True:
            # 保持连接，等待训练脚本推送的消息
            data = await websocket.receive_text()
            # 前端也可以发消息（比如取消训练）
            if data == "cancel":
                # TODO: 设置取消标志，训练脚本下次检查时退出
                pass
    except WebSocketDisconnect:
        manager.disconnect(websocket, task_id)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

// 前端 TrainingMonitor.vue — 训练监控组件
import { ref, onMounted, onUnmounted } from 'vue'
import * as echarts from 'echarts'

const props = defineProps({ taskId: String })

const status = ref('连接中...')
const lossData = ref([])    // [{epoch, loss, acc}, ...]
const chartRef = ref(null)
let chart = null
let ws = null

onMounted(() => {
  // 初始化ECharts图表
  chart = echarts.init(chartRef.value)
  chart.setOption({
    title: { text: '训练Loss曲线' },
    xAxis: { type: 'category', data: [] },
    yAxis: [
      { type: 'value', name: 'Loss' },
      { type: 'value', name: 'Accuracy', max: 1 }
    ],
    series: [
      { name: 'Loss', type: 'line', data: [], 
        smooth: true, itemStyle: { color: '#ff6b6b' } },
      { name: 'Accuracy', type: 'line', yAxisIndex: 1, 
        data: [], smooth: true, 
        itemStyle: { color: '#4ecdc4' } }
    ]
  })
  
  // 连接WebSocket
  ws = new WebSocket(`ws://localhost:8000/ws/${props.taskId}`)
  
  ws.onmessage = (event) => {
    const msg = JSON.parse(event.data)
    
    switch (msg.type) {
      case 'status':
        status.value = msg.status + ': ' + msg.msg
        break
      case 'metrics':
        // 追加数据点
        lossData.value.push({
          epoch: msg.epoch,
          loss: msg.loss,
          acc: msg.acc
        })
        // 更新图表
        chart.setOption({
          xAxis: { 
            data: lossData.value.map(d => `Epoch ${d.epoch}`) 
          },
          series: [
            { data: lossData.value.map(d => d.loss) },
            { data: lossData.value.map(d => d.acc) }
          ]
        })
        break
      case 'eval':
        // 显示评估结果
        status.value = `评估完成：准确率 ${(msg.accuracy * 100).toFixed(1)}%`
        break
      case 'done':
        status.value = '训练完成！'
        // 显示下载按钮
        downloadUrl.value = msg.download_url
        break
    }
  }
  
  ws.onclose = () => { status.value = '连接已断开' }
})

onUnmounted(() => {
  ws?.close()
  chart?.dispose()
})

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

# api/download.py — 部署包下载接口
from fastapi import APIRouter
from fastapi.responses import FileResponse

router = APIRouter(prefix="/api/download", tags=["download"])

@router.get("/{task_id}")
async def download_package(task_id: str):
    """下载训练完成的部署包
    
    返回的ZIP包含：
    ├── deploy_config.json    — 部署配置
    ├── model_int8.tflite     — 量化后的模型
    ├── mp_deployment_source/ — MicroPython推理脚本
    │   └── infer.py
    ├── eval_report.json      — 评估报告
    └── README.txt           — 部署说明
    """
    zip_path = f"/data/output/{task_id}/{task_id}_deploy.zip"
    
    if not Path(zip_path).exists():
        return {"error": "部署包不存在或任务未完成"}
    
    return FileResponse(
        zip_path,
        media_type="application/zip",
        filename=f"{task_id}_deploy.zip"
    )