Yolov8 Pruning

This repository is the original source. For a more detailed repository and accompanying video explanations【Chinese: shouxie_ai’s paid video tutorial, contact wechat: shouxie_ai; English: comming】

1. Pretrain[optional]
2. Constrained Training
3. Prune
4. Finetune

Takeaway: Experiments using the yolov8s model on VOC2007 showed pretraining and constrained training reaching a similar mAP50(B) of ~0.77 at epoch 50. After pruning, the finetuning phase took 65 epochs to achieve the same mAP50(B). The ONNX model size reduced from 43M to 36M after pruning.

1. Pretrain[optional]

git clone
use VOC2007, and modify the VOC.yaml
disable amp

# FILE: ultralytics/yolo/engine/trainer.py
...
def check_amp(model):
    # Avoid using mixed precision to affect finetune
    return False # <============== modified
    device = next(model.parameters()).device  # get model device
    if device.type in ('cpu', 'mps'):
        return False  # AMP only used on CUDA devices

    def amp_allclose(m, im):
        # All close FP32 vs AMP results
    ...

Training log

epoch	train/box_loss	train/cls_loss	train/dfl_loss	metrics/precision(B)	metrics/recall(B)	metrics/mAP50(B)	metrics/mAP50-95(B)	val/box_loss	val/cls_loss	val/dfl_loss	lr/pg0	lr/pg1	lr/pg2
0	0.92673	1.6385	1.2113	0.77207	0.77506	0.83161	0.63104	0.84319	0.8203	1.1008	0.070096	0.0033227	0.0033227
1	0.9417	1.1234	1.2114	0.67739	0.61248	0.66058	0.45227	1.0539	1.1543	1.2772	0.04003	0.0065902	0.0065902
2	1.0578	1.2969	1.296	0.58292	0.47764	0.50248	0.30112	1.2869	1.6265	1.4896	0.0098978	0.0097916	0.0097916
3	1.1906	1.5575	1.3958	0.54442	0.48757	0.48652	0.29632	1.3197	1.5872	1.5178	0.009703	0.009703	0.009703
4	1.204	1.5625	1.4142	0.57654	0.49493	0.51143	0.31174	1.2786	1.5057	1.4828	0.009703	0.009703	0.009703
5	1.1945	1.5082	1.4009	0.61871	0.52062	0.5703	0.35654	1.265	1.4636	1.4805	0.009604	0.009604	0.009604
6	1.1744	1.4468	1.3858	0.62485	0.57612	0.60489	0.3855	1.2203	1.3206	1.4285	0.009505	0.009505	0.009505
7	1.1517	1.4118	1.375	0.61661	0.59905	0.61871	0.39268	1.2074	1.3148	1.4195	0.009406	0.009406	0.009406
8	1.1299	1.3454	1.3518	0.66592	0.58576	0.62926	0.40537	1.1967	1.2811	1.4135	0.009307	0.009307	0.009307
9	1.1133	1.3284	1.348	0.67145	0.59822	0.64811	0.42689	1.1539	1.2208	1.3799	0.009208	0.009208	0.009208
10	1.1031	1.2827	1.3376	0.65911	0.61371	0.65124	0.42923	1.1579	1.2286	1.3815	0.009109	0.009109	0.009109
11	1.0827	1.2543	1.3268	0.71019	0.61398	0.68202	0.45785	1.1323	1.1531	1.3561	0.00901	0.00901	0.00901
12	1.0642	1.242	1.315	0.70052	0.61976	0.67969	0.45524	1.1335	1.154	1.3547	0.008911	0.008911	0.008911
13	1.0623	1.1975	1.3067	0.70224	0.6279	0.68803	0.46259	1.1211	1.1146	1.3568	0.008812	0.008812	0.008812
14	1.0489	1.1583	1.2964	0.71233	0.64626	0.70062	0.47288	1.1182	1.1015	1.3457	0.008713	0.008713	0.008713
15	1.0374	1.1518	1.2905	0.70464	0.63515	0.69193	0.46424	1.1156	1.0929	1.3454	0.008614	0.008614	0.008614
16	1.014	1.1291	1.2738	0.67909	0.64529	0.6864	0.46629	1.117	1.1117	1.3469	0.008515	0.008515	0.008515
17	1.0183	1.1054	1.2703	0.73264	0.64282	0.71155	0.48454	1.0994	1.0573	1.3332	0.008416	0.008416	0.008416
18	0.99619	1.0777	1.2573	0.72897	0.65939	0.71447	0.49011	1.0935	1.0545	1.3281	0.008317	0.008317	0.008317
19	0.98963	1.0468	1.2505	0.72029	0.66052	0.71277	0.49032	1.0875	1.0753	1.3206	0.008218	0.008218	0.008218
20	0.97962	1.0346	1.247	0.73563	0.6511	0.7142	0.49308	1.0702	1.0447	1.3112	0.008119	0.008119	0.008119
21	0.97159	1.0108	1.2417	0.72999	0.65638	0.71672	0.49605	1.073	1.0317	1.3106	0.00802	0.00802	0.00802
22	0.96612	0.99422	1.2354	0.73903	0.65636	0.71879	0.4968	1.0853	1.0214	1.3102	0.007921	0.007921	0.007921
23	0.95418	0.97995	1.2299	0.74019	0.6648	0.73817	0.51044	1.0775	1.0068	1.3188	0.007822	0.007822	0.007822
24	0.94679	0.96945	1.2227	0.73285	0.67695	0.72656	0.50312	1.0554	0.98838	1.2957	0.007723	0.007723	0.007723
25	0.93552	0.95821	1.217	0.74328	0.67199	0.72949	0.50698	1.0687	1.0016	1.3023	0.007624	0.007624	0.007624
26	0.93168	0.94349	1.2141	0.76862	0.66199	0.74154	0.5189	1.0617	0.98225	1.3112	0.007525	0.007525	0.007525
27	0.92401	0.93271	1.2139	0.7524	0.66939	0.73902	0.52226	1.0512	0.98525	1.297	0.007426	0.007426	0.007426
28	0.9235	0.90858	1.2047	0.74359	0.68755	0.73913	0.5187	1.0651	0.97496	1.3055	0.007327	0.007327	0.007327
29	0.91447	0.90524	1.2013	0.7685	0.6854	0.75282	0.53004	1.0501	0.94474	1.2915	0.007228	0.007228	0.007228
30	0.90837	0.88461	1.1886	0.74858	0.69356	0.74533	0.51988	1.0576	0.95429	1.2979	0.007129	0.007129	0.007129
31	0.89539	0.8767	1.1835	0.76213	0.67857	0.74776	0.52786	1.0522	0.95761	1.2928	0.00703	0.00703	0.00703
32	0.88759	0.8647	1.179	0.77072	0.68623	0.75324	0.53284	1.0454	0.95159	1.2888	0.006931	0.006931	0.006931
33	0.88108	0.85207	1.1745	0.76965	0.69816	0.76108	0.5428	1.0429	0.94663	1.2876	0.006832	0.006832	0.006832
34	0.88291	0.84293	1.1768	0.75702	0.70109	0.75217	0.53493	1.0437	0.92887	1.2811	0.006733	0.006733	0.006733
35	0.86054	0.83013	1.1648	0.77788	0.68466	0.75603	0.53917	1.0429	0.94194	1.2908	0.006634	0.006634	0.006634
36	0.85863	0.81798	1.1658	0.7657	0.69179	0.7546	0.53687	1.0328	0.93694	1.2801	0.006535	0.006535	0.006535
37	0.86479	0.80596	1.159	0.775	0.69038	0.7588	0.54172	1.0295	0.9394	1.2787	0.006436	0.006436	0.006436
38	0.84868	0.79285	1.1563	0.77168	0.69569	0.76378	0.54661	1.0268	0.90889	1.278	0.006337	0.006337	0.006337
40	0.83386	0.77434	1.1441	0.78373	0.68312	0.75814	0.54293	1.0315	0.91365	1.2843	0.006139	0.006139	0.006139
41	0.83605	0.76247	1.1441	0.76758	0.70613	0.76328	0.5461	1.0404	0.90907	1.286	0.00604	0.00604	0.00604
42	0.82537	0.75351	1.1383	0.77132	0.69906	0.75637	0.54226	1.0325	0.91358	1.2826	0.005941	0.005941	0.005941
43	0.82065	0.7466	1.1337	0.78819	0.69624	0.76861	0.5532	1.0267	0.90715	1.2777	0.005842	0.005842	0.005842
44	0.81619	0.74658	1.1335	0.76049	0.7084	0.76285	0.54782	1.0149	0.91284	1.2674	0.005743	0.005743	0.005743
45	0.80921	0.73341	1.1287	0.7871	0.70758	0.77009	0.55715	1.0166	0.89331	1.2704	0.005644	0.005644	0.005644
46	0.80473	0.73055	1.1264	0.78967	0.7021	0.77082	0.55884	1.0233	0.9001	1.2791	0.005545	0.005545	0.005545
47	0.79876	0.7237	1.1205	0.78231	0.71046	0.77262	0.55573	1.0241	0.88496	1.2789	0.005446	0.005446	0.005446
48	0.79629	0.7041	1.118	0.79295	0.69772	0.77068	0.55465	1.0314	0.88591	1.2802	0.005347	0.005347	0.005347
49	0.78996	0.70273	1.1133	0.79776	0.70282	0.77587	0.56048	1.0228	0.87524	1.2721	0.005248	0.005248	0.005248
50	0.78926	0.70901	1.1161	0.80225	0.69859	0.773	0.56149	1.0208	0.89449	1.2797	0.005149	0.005149	0.005149

2. Constrained Training

prune the BN layer by adding L1 regularizer.

# FILE: ultralytics/yolo/engine/trainer.py
...
# Backward
self.scaler.scale(self.loss).backward()

# <============ added
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
if ni - last_opt_step >= self.accumulate:
    self.optimizer_step()
    last_opt_step = ni
...

Training log

epoch	train/box_loss	train/cls_loss	train/dfl_loss	metrics/precision(B)	metrics/recall(B)	metrics/mAP50(B)	metrics/mAP50-95(B)	val/box_loss	val/cls_loss	val/dfl_loss	lr/pg0	lr/pg1	lr/pg2
0	0.92565	1.6423	1.2109	0.78443	0.76733	0.82516	0.62045	0.86821	0.84614	1.1113	0.070096	0.0033227	0.0033227
1	0.94448	1.1347	1.2122	0.73125	0.67825	0.73839	0.52702	0.96679	1.0091	1.2034	0.04003	0.0065902	0.0065902
2	1.0604	1.3009	1.2986	0.6136	0.54731	0.57973	0.37943	1.1684	1.4044	1.3756	0.0098978	0.0097916	0.0097916
3	1.196	1.5631	1.3974	0.58267	0.48531	0.51603	0.31687	1.2998	1.5516	1.5155	0.009703	0.009703	0.009703
4	1.2029	1.5454	1.4114	0.6036	0.5207	0.55252	0.34034	1.2789	1.4587	1.4855	0.009703	0.009703	0.009703
5	1.202	1.5032	1.4049	0.5087	0.40369	0.40513	0.23736	1.4244	1.8939	1.6476	0.009604	0.009604	0.009604
6	1.1764	1.4465	1.3887	0.63867	0.55878	0.60577	0.38883	1.2146	1.2978	1.4348	0.009505	0.009505	0.009505
7	1.1522	1.3962	1.3754	0.656	0.5712	0.62207	0.40084	1.237	1.3347	1.4478	0.009406	0.009406	0.009406
8	1.127	1.3373	1.354	0.67759	0.58	0.63802	0.41447	1.1818	1.2604	1.3894	0.009307	0.009307	0.009307
9	1.1105	1.3162	1.3473	0.64599	0.56469	0.61175	0.39579	1.1895	1.2821	1.4132	0.009208	0.009208	0.009208
10	1.1008	1.2829	1.3399	0.69645	0.61177	0.66806	0.44422	1.1531	1.1822	1.3768	0.009109	0.009109	0.009109
11	1.0765	1.2522	1.3232	0.71631	0.60926	0.68153	0.4529	1.1446	1.136	1.3674	0.00901	0.00901	0.00901
12	1.0662	1.2235	1.3134	0.69318	0.62751	0.67452	0.45073	1.132	1.1636	1.3594	0.008911	0.008911	0.008911
13	1.0616	1.1851	1.2999	0.71751	0.61549	0.68857	0.46305	1.1255	1.1331	1.347	0.008812	0.008812	0.008812
14	1.0494	1.1522	1.2947	0.69513	0.63001	0.67925	0.45575	1.12	1.1272	1.3465	0.008713	0.008713	0.008713
15	1.037	1.1478	1.2873	0.70825	0.63711	0.68977	0.46421	1.1219	1.097	1.3403	0.008614	0.008614	0.008614
16	1.0175	1.1304	1.2747	0.73014	0.6315	0.70322	0.47744	1.1077	1.077	1.3344	0.008515	0.008515	0.008515
17	1.0172	1.0993	1.269	0.74106	0.65394	0.72733	0.49715	1.0884	1.0431	1.3239	0.008416	0.008416	0.008416
18	1.0009	1.0633	1.2578	0.73366	0.6424	0.71016	0.4857	1.0928	1.0509	1.3175	0.008317	0.008317	0.008317
19	0.99458	1.0608	1.2552	0.73498	0.65621	0.72397	0.50389	1.0748	1.0456	1.3113	0.008218	0.008218	0.008218
20	0.97939	1.0313	1.2465	0.75441	0.64815	0.7261	0.50359	1.0737	1.0337	1.3059	0.008119	0.008119	0.008119
21	0.97806	1.026	1.2422	0.72394	0.63557	0.70287	0.48524	1.0717	1.0369	1.2972	0.00802	0.00802	0.00802
22	0.97879	1.0127	1.2431	0.7394	0.65045	0.71215	0.48902	1.0888	1.0292	1.3174	0.007921	0.007921	0.007921
23	0.95895	0.99651	1.232	0.73731	0.67343	0.73924	0.51158	1.0691	0.99626	1.3042	0.007822	0.007822	0.007822
24	0.95615	0.99158	1.2322	0.73624	0.66665	0.72512	0.50267	1.0644	1.0126	1.3008	0.007723	0.007723	0.007723
25	0.94517	0.97143	1.2231	0.75806	0.65886	0.73128	0.50956	1.0688	0.99726	1.3039	0.007624	0.007624	0.007624
26	0.93639	0.95727	1.2188	0.74016	0.67544	0.73665	0.5112	1.0563	0.98787	1.3013	0.007525	0.007525	0.007525
27	0.93284	0.94819	1.2209	0.73318	0.67589	0.73246	0.51282	1.0617	0.9874	1.2937	0.007426	0.007426	0.007426
28	0.9343	0.93543	1.2115	0.74821	0.68133	0.74574	0.52401	1.0533	0.97285	1.2896	0.007327	0.007327	0.007327
29	0.92951	0.93541	1.2125	0.75947	0.67266	0.74126	0.51764	1.054	0.96946	1.2953	0.007228	0.007228	0.007228
30	0.92688	0.91381	1.2025	0.74962	0.67131	0.73941	0.5137	1.0758	0.9762	1.3053	0.007129	0.007129	0.007129
31	0.91471	0.91552	1.1992	0.75604	0.67727	0.7467	0.52798	1.0448	0.96641	1.2824	0.00703	0.00703	0.00703
32	0.90619	0.89194	1.1904	0.74516	0.68767	0.75086	0.52798	1.0448	0.95155	1.2845	0.006931	0.006931	0.006931
33	0.9009	0.89107	1.1897	0.74989	0.67013	0.74297	0.52946	1.0401	0.96254	1.287	0.006832	0.006832	0.006832
34	0.903	0.88358	1.1913	0.75685	0.68489	0.74407	0.52314	1.053	0.95245	1.2863	0.006733	0.006733	0.006733
35	0.88498	0.87044	1.1804	0.76152	0.6916	0.75495	0.53359	1.046	0.94952	1.2841	0.006634	0.006634	0.006634
36	0.88366	0.86359	1.1834	0.75568	0.70024	0.75525	0.53069	1.0468	0.9341	1.2886	0.006535	0.006535	0.006535
37	0.88893	0.85432	1.1778	0.76405	0.68825	0.75449	0.53531	1.0425	0.94635	1.2855	0.006436	0.006436	0.006436
38	0.87813	0.85042	1.178	0.76257	0.67111	0.74801	0.52377	1.0624	0.9561	1.3098	0.006337	0.006337	0.006337
39	0.88227	0.84591	1.172	0.77845	0.69019	0.7627	0.53847	1.0465	0.92067	1.2847	0.006238	0.006238	0.006238
40	0.87228	0.82927	1.1709	0.77537	0.69448	0.76192	0.5392	1.0454	0.91643	1.2826	0.006139	0.006139	0.006139
41	0.86389	0.82287	1.1627	0.7626	0.7069	0.76237	0.54594	1.0379	0.91307	1.2755	0.00604	0.00604	0.00604
42	0.85961	0.81907	1.161	0.77033	0.70088	0.76257	0.54426	1.0371	0.90801	1.2758	0.005941	0.005941	0.005941
43	0.85414	0.80214	1.1564	0.76815	0.69915	0.76696	0.54287	1.0426	0.9051	1.281	0.005842	0.005842	0.005842
44	0.85261	0.81496	1.1614	0.77	0.69292	0.76054	0.53674	1.0424	0.92179	1.277	0.005743	0.005743	0.005743
45	0.85068	0.8013	1.1564	0.78907	0.69016	0.77027	0.54926	1.0389	0.90745	1.2772	0.005644	0.005644	0.005644
46	0.84368	0.80178	1.1515	0.79301	0.69554	0.7723	0.55095	1.0459	0.90891	1.2887	0.005545	0.005545	0.005545
47	0.84028	0.79458	1.1483	0.76734	0.70468	0.76809	0.54484	1.044	0.89653	1.283	0.005446	0.005446	0.005446
48	0.83392	0.77818	1.1443	0.77305	0.69609	0.76721	0.54672	1.036	0.89301	1.2787	0.005347	0.005347	0.005347
49	0.83255	0.78041	1.1439	0.76445	0.71068	0.76725	0.54742	1.0361	0.89842	1.2748	0.005248	0.005248	0.005248
50	0.829	0.76955	1.1423	0.7801	0.70469	0.77336	0.55284	1.0301	0.88995	1.2764	0.005149	0.005149	0.005149

3. Prune

3.1 Inspect the Biases of BN

After pruning, ensure most of the bias of BN are small enought to close to 0. Otherwise, go back to constrained training

for name, m in model.named_modules():
    if isinstance(m, torch.nn.BatchNorm2d):
        w = m.weight.abs().detach()
        b = m.bias.abs().detach()
        ws.append(w)
        bs.append(b)
        print(name, w.max().item(), w.min().item(), b.max().item(), b.min().item())

3.2 Set Threshold and Pruning Rate

threshold: global or local
factor: the ratio of keeping. Pruning too much is not recommended.

factor = 0.8
ws = torch.cat(ws)
threshold = torch.sort(ws, descending=True)[0][int(len(ws) * factor)]
print(threshold)

3.3 TB-pattern pruning(Top-Bottom)

TB-pattens
- naive 1 vs 1 : Bottleneck;
- Seq: C2f and SPPF
- Detect: FPN

Related Module Code

Show / hide Bottleneck code

# bottleneck.png
class Bottleneck(nn.Module):
    # Standard bottleneck
    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):  # ch_in, ch_out, shortcut, groups, kernels, expand
        super().__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, k[0], 1)
        self.cv2 = Conv(c_, c2, k[1], 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))

Show / hide C2f code

# c2f.png
class C2f(nn.Module):
    # CSP Bottleneck with 2 convolutions
    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super().__init__()
        self.c = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

    def forward(self, x):
        y = list(self.cv1(x).chunk(2,1))
        y.extend(m(y[-1]) for m in self.m) # y[-1] is the right branch
        return self.cv2(torch.cat(y, 1))

Show / hide SPPF code

# sppf.png
class SPPF(nn.Module):
    # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
    def __init__(self, c1, c2, k=5):  # equivalent to SPP(k=(5, 9, 13))
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * 4, c2, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

    def forward(self, x):
        x = self.cv1(x)
        y1 = self.m(x)
        y2 = self.m(y1)
        return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))

4. Finetune

diable l1 constraint
avoid resume from yaml. Refer to 4.7补充细节_剪枝后不要从yaml导入结构.mp4

# FILE: ultralytics/yolo/engine/model.py
# self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)
# self.model = self.trainer.model
self.trainer.model = self.model.train()

ensure all parameters require grad

Name		Name	Last commit message	Last commit date
Latest commit History 6 Commits
docker		docker
docs		docs
examples		examples
fig		fig
tests		tests
ultralytics		ultralytics
wandb		wandb
CITATION.cff		CITATION.cff
CONTRIBUTING.md		CONTRIBUTING.md
LICENSE		LICENSE
MANIFEST.in		MANIFEST.in
README.md		README.md
README.zh-CN.md		README.zh-CN.md
excluded.txt		excluded.txt
export.py		export.py
mkdocs.yml		mkdocs.yml
prune.py		prune.py
prune_by_su.py		prune_by_su.py
requirements.txt		requirements.txt
setup.cfg		setup.cfg
setup.py		setup.py
train.py		train.py
yolov8README.md		yolov8README.md
yolov8_for_reference.yaml		yolov8_for_reference.yaml
yolov8s.pt		yolov8s.pt

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Repository files navigation

Yolov8 Pruning

1. Pretrain[optional]

2. Constrained Training

3. Prune

3.1 Inspect the Biases of BN

3.2 Set Threshold and Pruning Rate

3.3 TB-pattern pruning(Top-Bottom)

4. Finetune

About

Releases

Packages

Languages

License

downpeopleup/yolov8-pruning

Folders and files

Latest commit

History

Repository files navigation

Yolov8 Pruning

1. Pretrain[optional]

2. Constrained Training

3. Prune

3.1 Inspect the Biases of BN

3.2 Set Threshold and Pruning Rate

3.3 TB-pattern pruning(Top-Bottom)

4. Finetune

About

Resources

License

Security policy

Stars

Watchers

Forks

Releases

Packages 0

Languages

Packages