yolov11剪枝、蒸馏、加注意力

这段代码是一个用于深度学习模型训练的Python脚本,特别是针对目标检测任务,使用了YOLO(You Only Look Once)算法。代码中包含了多个步骤,每个步骤都是模型训练过程中的一个阶段。以下是对代码的详细解释:

完整版代码在GitHub上:yolov11剪枝蒸馏

1. **导入必要的库和模块**:

  • `from ultralytics import YOLO`:导入了ultralytics提供的YOLO模型库。

  • `import os`:导入操作系统接口模块,用于文件和目录操作。

  • `from utils.yolo.attention import add_attention`:导入一个自定义模块,用于给模型添加注意力机制。

2. **设置环境变量和路径**:

  • `os.environ["CUDA_VISIBLE_DEVICES"]="0,1"`:这行代码被注释掉了,它的作用是设置CUDA环境变量,指定使用哪几个GPU设备。在这里指定了0号和1号GPU。

  • `root = os.getcwd()`:获取当前工作目录的路径。

  • `name_yaml`、`name_pretrain`等变量定义了配置文件和预训练模型文件的路径。

3. **定义训练步骤**:

  • `step1_train()`:加载预训练模型并开始训练。

  • `step2_Constraint_train()`:在约束条件下进行训练,例如可能涉及到正则化或其他约束条件。

  • `step3_pruning()`:使用自定义的`do_pruning`函数对模型进行剪枝,以减少模型的复杂度。

  • `step4_finetune()`:微调剪枝后的模型。

  • `step5_distillation()`:使用知识蒸馏技术,将一个训练好的大模型(教师模型)的知识传递给一个较小的模型(学生模型)。

4. **训练函数参数解释**:

  • `data`:指定数据配置文件的路径。

  • `device`:指定训练使用的设备,如GPU。

  • `imgsz`:指定输入图像的大小。

  • `epochs`:指定训练的轮数。

  • `batch`:指定每批训练的样本数量。

  • `workers`:指定用于数据加载的工作线程数量。

  • `save_period`:指定保存模型的周期。

  • `name`:指定模型保存的路径。

  • `amp`:指定是否使用自动混合精度训练。

  • `Distillation`:指定知识蒸馏的教师模型。

  • `loss_type`:指定损失函数的类型。

  • `layers`:指定进行蒸馏的层。

python 复制代码
from ultralytics import YOLO
import os
from utils.yolo.attention import add_attention
# os.environ["CUDA_VISIBLE_DEVICES"]="0,1"

root = os.getcwd()
## 配置文件路径
name_yaml = os.path.join(root, "data.yaml")
name_pretrain = os.path.join(root, "runs/segment/ori/weights/best.pt")
## 原始训练路径
# path_train = os.path.join(root, "runs/detect/VOC")
name_train = "runs/segment/ori/weights/best.pt"
## 约束训练路径、剪枝模型文件
path_constraint_train = os.path.join(root, "runs/segment/Constraint")
name_prune_before = os.path.join(path_constraint_train, "weights/last.pt")
name_prune_after = os.path.join(path_constraint_train, "weights/prune.pt")
## 微调路径
path_fineturn = os.path.join(root, "runs/detect/VOC_finetune")

def step1_train():
    model = YOLO(name_pretrain)
    model.train(data=name_yaml, device="0", imgsz=720, epochs=50, batch=2, workers=0, save_period=1)  # train the model


## 2024.3.4添加【amp=False】
def step2_Constraint_train():
    model = YOLO(name_train)
    model.train(data=name_yaml, device="0", imgsz=640, epochs=50, batch=2, amp=False, workers=0, save_period=1,
                name=path_constraint_train)  # train the model


def step3_pruning():
    from utils.yolo.LL_pruning import do_pruning
    do_pruning(name_prune_before, name_prune_after)


def step4_finetune():
    model = YOLO(name_prune_after)  # load a pretrained model (recommended for training)
    for param in model.parameters():
        param.requires_grad = True
    model.train(data=name_yaml, device="0", imgsz=640, epochs=200, batch=2, workers=0, name=path_fineturn)  # train the model

def step5_distillation():
    layers = ["6", "8", "13", "16", "19", "22"]
    model_t = YOLO('runs/segment/ori/weights/best.pt')  # the teacher model
    model_s = YOLO('runs/segment/Constraint/weights/prune.pt')  # the student model
    model_s = add_attention(model_s)
    """
    Attributes:
        Distillation: the distillation model
    """
    model_s.train(data="data.yaml", Distillation=model_t.model, loss_type='mgd',layers=layers, amp=False, imgsz=1280, epochs=300,
                  batch=2, device=0, workers=0, lr0=0.001)


if __name__ == '__main__':
    # step1_train()
    # step2_Constraint_train()
    # step3_pruning()
    # step4_finetune()
    step5_distillation()
相关推荐
阿坡RPA2 小时前
手搓MCP客户端&服务端:从零到实战极速了解MCP是什么?
人工智能·aigc
用户27784491049932 小时前
借助DeepSeek智能生成测试用例:从提示词到Excel表格的全流程实践
人工智能·python
机器之心2 小时前
刚刚,DeepSeek公布推理时Scaling新论文,R2要来了?
人工智能
算AI4 小时前
人工智能+牙科:临床应用中的几个问题
人工智能·算法
凯子坚持 c5 小时前
基于飞桨框架3.0本地DeepSeek-R1蒸馏版部署实战
人工智能·paddlepaddle
你觉得2055 小时前
哈尔滨工业大学DeepSeek公开课:探索大模型原理、技术与应用从GPT到DeepSeek|附视频与讲义下载方法
大数据·人工智能·python·gpt·学习·机器学习·aigc
8K超高清6 小时前
中国8K摄像机:科技赋能文化传承新图景
大数据·人工智能·科技·物联网·智能硬件
hyshhhh6 小时前
【算法岗面试题】深度学习中如何防止过拟合?
网络·人工智能·深度学习·神经网络·算法·计算机视觉
薛定谔的猫-菜鸟程序员6 小时前
零基础玩转深度神经网络大模型:从Hello World到AI炼金术-详解版(含:Conda 全面使用指南)
人工智能·神经网络·dnn
币之互联万物6 小时前
2025 AI智能数字农业研讨会在苏州启幕,科技助农与数据兴业成焦点
人工智能·科技