当我们使用yolov8训练时候,保存的模型变大,如下图:
原模型
训练出来的模型
经过仔细调查,发现是保存的模型中多了很多数据。
原模型
训练出来的模型
只需要把文件中.pt文件读取,重写一遍保存。
from ultralytics import YOLO
import torch
from ultralytics.nn.modules import Bottleneck, Conv, C2f, SPPF, Detect
from copy import deepcopy
# Load a model
yolo = YOLO("./runs/detect/train2/weights/last.pt")
ckpt = {
'epoch': -1,
'best_fitness': None,
'model': yolo.ckpt["model"],
'ema': yolo.ckpt['ema'],
'updates': None,
'optimizer': None,
'train_args': yolo.ckpt["train_args"], # save as dict
'date': None,
'version': '8.0.142'}
torch.save(ckpt, "./runs/detect/train1/weights/new.pt")
如果把ema也设置成None,就文件跟原模型一样大小。
在训练过程中,EMA权重会逐渐接近模型权重,但以一种缓慢的方式,因为EMA权重的更新会考虑到之前的EMA权重。EMA可以帮助模型从训练过程中的局部最小值中恢复,并且通常可以在训练结束时提供更稳定的模型参数。
在实际应用中,EMA模型的权重可能不会用于模型的训练更新,而是作为辅助,用于在训练过程中或训练结束后生成更稳健的模型参数。有时,训练完成后,可以使用EMA权重作为模型的最终参数,因为它们可能提供了更好的泛化性能。