1. 模型稀疏化训练(增加权重稀疏化的代码进行训练)
--ultralytics/engine/trainer.py
约在345行:
# Backward
self.scaler.scale(self.loss).backward()
# ========== 新增 ==========
l1_lambda = 1e-2 * (1 - 0.9 * epoch / self.epochs)
for k, m in self.model.named_modules():
if isinstance(m, nn.BatchNorm2d):
m.weight.grad.data.add_(l1_lambda * torch.sign(m.weight.data))
m.bias.grad.data.add_(1e-2 * torch.sign(m.bias.data))
# ========== 新增 ==========
# Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
2. 开始训练
--ultralytics-main/train.py
训练代码如下:
from ultralytics import YOLO
import torch
model = YOLO('yolov8s.yaml'