超级会员免费看
本文介绍了如何对YOLOv8模型进行稀疏化训练和剪枝，包括在训练过程中引入权重稀疏化的代码，通过`ultralytics`库进行模型的训练和剪枝，以及剪枝后的模型微调步骤。注意在剪枝和再次训练时对代码的特定修改。
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.htm
订阅专栏 解锁全文
扫一扫
1056
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
