When training a Yolo model like below:
from ultralytics import YOLO
# Load a model
model = YOLO('yolov8s.pt')
results = model.train(data='coco128.yaml',
epochs=100, imgsz=640, save_period=1)
The save_period option will save every epoch. When the best epoch is found the file is saved as best.pt.
The file size of best.pt is ~27MB and each epoch is ~120MB. Is it possible to use the compression applied to the best epoch at the end of training to every epoch (even if this is done after training).
What you're seeing has nothing to do with compression. When saving a checkpoint during training, the YOLO training script saves (in addition to other metadata) weights of the current model, weights of the models EMA, and the optimizers parameters. These are needed to continue training from a checkpoint, but only the EMA is kept for inference.
To remove the optimizer and current model, you can use ultralytics' utility function strip_optimizer
, which strips the unnecessary weights from the state dict found at a given path, and saves the smaller file at the same location.
from ultralytics.utils.torch_utils import strip_optimizer
path_to_checkpoint = "best.pt"
strip_optimizer(path_to_checkpoint)