The size of YOLOv5s is 14 megabytes

[leeyunhome]

❔Question

The size of YOLOv5s is 14 megabytes, but when I actually train it, it comes out to 15 megs.
Could there be such a difference?

Additional context

[glenn-jocher]

@leeyunhome the size of the latest YOLOv5 models is shown in the assets section of the releases. These models were trained using the same exact steps as in our Train Custom Data tutorial:
https://docs.ultralytics.com/yolov5/tutorials/train_custom_data

https://github.com/ultralytics/yolov5/tree/v5.0

[github-actions]

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

• Wiki – https://github.com/ultralytics/yolov5/wiki
• Tutorials – https://docs.ultralytics.com/yolov5
• Docs – https://docs.ultralytics.com

Access additional Ultralytics ⚡ resources:

• Ultralytics HUB – https://ultralytics.com
• Vision API – https://ultralytics.com/yolov5
• About Us – https://ultralytics.com/about
• Join Our Team – https://ultralytics.com/work
• Contact Us – https://ultralytics.com/contact

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐!

[dhaval-zala-aivid]

@glenn-jocher
I'm training yolov5n model. Ideally it should generate 4mb file of best.pt and last.pt. But it 14.7 mb file.
here is the training logs:

python train.py --img 640 --batch 8 --epochs 50 --data data/coco.yaml --weights yolov5n.pt --exist-ok --name /home/worker2/Dhavalsinh/HeadDetection/YOLOV5/yolov5/runs
train: weights=yolov5n.pt, cfg=, data=data/coco.yaml, hyp=data/hyps/hyp.scratch-low.yaml, epochs=50, batch_size=8, imgsz=640, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, noplots=False, evolve=None, bucket=, cache=None, image_weights=False, device=, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=runs/train, name=/home/worker2/Dhavalsinh/HeadDetection/YOLOV5/yolov5/runs, exist_ok=True, quad=False, cos_lr=False, label_smoothing=0.0, patience=100, freeze=[0], save_period=-1, seed=0, local_rank=-1, entity=None, upload_dataset=False, bbox_interval=-1, artifact_alias=latest
error: cannot lock ref 'refs/remotes/origin/docker': 'refs/remotes/origin/docker/20_08' exists; cannot create 'refs/remotes/origin/docker'
From https://github.com/ultralytics/yolov5
 ! [new branch]        docker     -> origin/docker  (unable to update local ref)
Command 'git fetch origin' returned non-zero exit status 1.
YOLOv5 🚀 v6.2-224-g82a55855 Python-3.7.10 torch-1.10.0+cu102 CUDA:0 (NVIDIA GeForce GTX 1660 SUPER, 5942MiB)

hyperparameters: lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=0.05, cls=0.5, cls_pw=1.0, obj=1.0, obj_pw=1.0, iou_t=0.2, anchor_t=4.0, fl_gamma=0.0, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0
ClearML: run 'pip install clearml' to automatically track, visualize and remotely train YOLOv5 🚀 in ClearML
Comet: run 'pip install comet_ml' to automatically track and visualize YOLOv5 🚀 runs in Comet
TensorBoard: Start with 'tensorboard --logdir /home/worker2/Dhavalsinh/HeadDetection/YOLOV5/yolov5', view at http://localhost:6006/
Overriding model.yaml nc=80 with nc=3

                 from  n    params  module                                  arguments                     
  0                -1  1      1760  models.common.Conv                      [3, 16, 6, 2, 2]              
  1                -1  1      4672  models.common.Conv                      [16, 32, 3, 2]                
  2                -1  1      4800  models.common.C3                        [32, 32, 1]                   
  3                -1  1     18560  models.common.Conv                      [32, 64, 3, 2]                
  4                -1  2     29184  models.common.C3                        [64, 64, 2]                   
  5                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               
  6                -1  3    156928  models.common.C3                        [128, 128, 3]                 
  7                -1  1    295424  models.common.Conv                      [128, 256, 3, 2]              
  8                -1  1    296448  models.common.C3                        [256, 256, 1]                 
  9                -1  1    164608  models.common.SPPF                      [256, 256, 5]                 
 10                -1  1     33024  models.common.Conv                      [256, 128, 1, 1]              
 11                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 12           [-1, 6]  1         0  models.common.Concat                    [1]                           
 13                -1  1     90880  models.common.C3                        [256, 128, 1, False]          
 14                -1  1      8320  models.common.Conv                      [128, 64, 1, 1]               
 15                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 16           [-1, 4]  1         0  models.common.Concat                    [1]                           
 17                -1  1     22912  models.common.C3                        [128, 64, 1, False]           
 18                -1  1     36992  models.common.Conv                      [64, 64, 3, 2]                
 19          [-1, 14]  1         0  models.common.Concat                    [1]                           
 20                -1  1     74496  models.common.C3                        [128, 128, 1, False]          
 21                -1  1    147712  models.common.Conv                      [128, 128, 3, 2]              
 22          [-1, 10]  1         0  models.common.Concat                    [1]                           
 23                -1  1    296448  models.common.C3                        [256, 256, 1, False]          
 24      [17, 20, 23]  1     10824  models.yolo.Detect                      [3, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [64, 128, 256]]
Model summary: 214 layers, 1767976 parameters, 1767976 gradients, 4.2 GFLOPs

What could be the reason and how can I solve it?

[glenn-jocher]

👋 Hello! Thanks for asking about training checkpoints. YOLOv5 🚀 checkpoints should be about 4X the size of final trained checkpoints, as they carry not just a FP16 model, but a FP16 EMA and an FP32 optimizer of the same size as the model (each model parameter has its own FP32 gradient saved within the optimizer). Final checkpoints contain only an FP16 model, with the EMA and optimizer both stripped after the final epoch of training:

yolov5/train.py

Lines 423 to 425 in b4a29b5

	for f in last, best:
	if f.exists():
	strip_optimizer(f) # strip optimizers

The strip_optimizer() function updates the checkpoint dictionary value for model, replacing it with ema, and sets ema and optimizer keys to None, which will reduce the checkpoint size by 3/4.

yolov5/utils/general.py

Lines 741 to 755 in b4a29b5

	def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer()
	# Strip optimizer from 'f' to finalize training, optionally save as 's'
	x = torch.load(f, map_location=torch.device('cpu'))
	if x.get('ema'):
	x['model'] = x['ema'] # replace model with ema
	for k in 'optimizer', 'best_fitness', 'wandb_id', 'ema', 'updates': # keys
	x[k] = None
	x['epoch'] = -1
	x['model'].half() # to FP16
	for p in x['model'].parameters():
	p.requires_grad = False
	torch.save(x, s or f)
	mb = os.path.getsize(s or f) / 1E6 # filesize
	print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")

Final trained checkpoints after strip_optimizer() should match the README table sizes below.

You can also run strip_optimizer() manually on any checkpoint to convert it into finalized weights ready to train a new model:

from utils.general import strip_optimizer

strip_optimizer('path/to/best.pt')
!python train.py --weights path/to/best.pt  # use best.pt to train a new model

Model Exports

Exported models may be larger or smaller than their PyTorch starting points depending on their quantization. i.e. FP16 models will be the same size, FP32 models will be 2X the size, and INT8 models will be 0.5X the size. See YOLOv5 Export tutorial for quantization options on model export.

Good luck 🍀 and let us know if you have any other questions!