YOLOv9模型重新参数化,将yolo.pt转为yolo-converted.pt
1、yolo-m转为yolo-m-converted
在我们使用train_dual.py训练好自己的模型后,我们会得到yolo系列的模型,比如我使用yolov9-m在数据集CityScapes上训练后,得到cityscapes-yolov9-m.pt,然后运行yolov9-m_to_converted.py代码得到cityscapes-yolov9-m-converted.pt,代码如下:记得修改nc为自己数据集的分类数,然后修改自己的加载模型路径和模型保存路径。
#!/usr/bin/env python
# coding: utf-8import torch
from models.yolo import Model# ## Convert YOLOv9-M# In[ ]:device = torch.device("cpu")
cfg = "../models/detect/gelan-m.yaml"
model = Model(cfg, ch=3, nc=8, anchors=3)
#model = model.half()
model = model.to(device)
_ = model.eval()
ckpt = torch.load('../runs/train/YOLOv9-m训练Cityscapes/weights/cityscapes-yolov9-m.pt', map_location='cpu')
model.names = ckpt['model'].names
model.nc = ckpt['model'].nc# In[ ]:idx = 0
for k, v in model.state_dict().items():if "model.{}.".format(idx) in k:if idx < 22:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx+1))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv2.".format(idx) in k:kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv3.".format(idx) in k:kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.dfl.".format(idx) in k:kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")else:while True:idx += 1if "model.{}.".format(idx) in k:breakif idx < 22:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx+1))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv2.".format(idx) in k:kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv3.".format(idx) in k:kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.dfl.".format(idx) in k:kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")
_ = model.eval()# In[ ]:m_ckpt = {'model': model.half(),'optimizer': None,'best_fitness': None,'ema': None,'updates': None,'opt': None,'git': None,'date': None,'epoch': -1}
torch.save(m_ckpt, "./cityscapes-yolov9-m-converted.pt")
运行效果如下:
2、 使用重参后的模型进行推理
得到cityscapes-yolov9-m-converted.pt后,权重大小为38.7MB,使用detect.py进行检测
可以看到模型信息为:377 layers, 19926616 parameters, 9048 gradients, 76.0 GFLOPs
此时在使用原来的cityscapes-yolov9-m.pt,权重大小为63.1MB,进行检测,效果如下:
可以看到模型信息为:588 layers, 32563288 parameters, 0 gradients, 130.7 GFLOPs
对于示例图片,两个模型都做到了一致的检测结果:7 cars, 2 persons, 1 rider, 1 bicycle,这检测效果还是很好的,很多远一点、小一点的目标我自己肉眼都没看到。
无论是推理速度还是模型参数,经过重参后都有所提升,并且官方称两者的精度是一致的,做到了既保留了较高的精度,又缩小了模型复杂度。
但是!虽然经过重参后的模型能够使用detect.py进行检测推理,目前还无法使用val或者val_dual.py对重新参数化的模型进行验证,比如cityscapes-yolov9-m.pt经过参数重新化得到cityscapes-yolov9-m-converted.pt,将converted模型用于验证会报错,只能等后续官方维护了。
3、 yolo-s转为yolo-s-converted
运行yolov9-s_to_converted.py代码得到yolov9-s-converted.pt,代码如下:记得修改nc为自己数据集的分类数,然后修改自己的加载模型路径和模型保存路径。
#!/usr/bin/env python
# coding: utf-8import torch
from models.yolo import Model# ## Convert YOLOv9-S# In[ ]:device = torch.device("cpu")
cfg = "../models/detect/gelan-s.yaml"
model = Model(cfg, ch=3, nc=80, anchors=3)
#model = model.half()
model = model.to(device)
_ = model.eval()
ckpt = torch.load('../models/pretrained-weights/yolov9-s.pt', map_location='cpu')
model.names = ckpt['model'].names
model.nc = ckpt['model'].nc# In[ ]:idx = 0
for k, v in model.state_dict().items():if "model.{}.".format(idx) in k:if idx < 22:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv2.".format(idx) in k:kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv3.".format(idx) in k:kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.dfl.".format(idx) in k:kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")else:while True:idx += 1if "model.{}.".format(idx) in k:breakif idx < 22:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv2.".format(idx) in k:kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv3.".format(idx) in k:kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.dfl.".format(idx) in k:kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")
_ = model.eval()# In[ ]:m_ckpt = {'model': model.half(),'optimizer': None,'best_fitness': None,'ema': None,'updates': None,'opt': None,'git': None,'date': None,'epoch': -1}
torch.save(m_ckpt, "./test_yolov9-s-converted.pt")
4、 yolo-c转为yolo-c-converted
运行yolov9-c_to_converted.py代码得到yolov9-c-converted.pt,代码如下:记得修改nc为自己数据集的分类数,然后修改自己的加载模型路径和模型保存路径。
#!/usr/bin/env python
# coding: utf-8import torch
from models.yolo import Model# ## Convert YOLOv9-C# In[ ]:device = torch.device("cpu")
cfg = "../models/detect/gelan-c.yaml"
model = Model(cfg, ch=3, nc=8, anchors=3)
#model = model.half()
model = model.to(device)
_ = model.eval()
ckpt = torch.load('../runs/train/YOLOv9-m训练Cityscapes/weights/best.pt', map_location='cpu')
model.names = ckpt['model'].names
model.nc = ckpt['model'].nc# In[ ]:idx = 0
for k, v in model.state_dict().items():if "model.{}.".format(idx) in k:if idx < 22:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx+1))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]elif "model.{}.cv2.".format(idx) in k:kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]elif "model.{}.cv3.".format(idx) in k:kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]elif "model.{}.dfl.".format(idx) in k:kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]else:while True:idx += 1if "model.{}.".format(idx) in k:breakif idx < 22:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx+1))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]elif "model.{}.cv2.".format(idx) in k:kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]elif "model.{}.cv3.".format(idx) in k:kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]elif "model.{}.dfl.".format(idx) in k:kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+16))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]
_ = model.eval()# In[ ]:m_ckpt = {'model': model.half(),'optimizer': None,'best_fitness': None,'ema': None,'updates': None,'opt': None,'git': None,'date': None,'epoch': -1}
torch.save(m_ckpt, "./cityscapes-yolov9-c-converted.pt")
5、 yolo-e转为yolo-e-converted
运行yolov9-e_to_converted.py代码得到yolov9-e-converted.pt,代码如下:记得修改nc为自己数据集的分类数,然后修改自己的加载模型路径和模型保存路径。
#!/usr/bin/env python
# coding: utf-8import torch
from models.yolo import Model# ## Convert YOLOv9-E# In[ ]:device = torch.device("cpu")
cfg = "../models/detect/gelan-e.yaml"
model = Model(cfg, ch=3, nc=8, anchors=3)
#model = model.half()
model = model.to(device)
_ = model.eval()
ckpt = torch.load('../runs/train/YOLOv9-m训练Cityscapes/weights/best.pt', map_location='cpu')
model.names = ckpt['model'].names
model.nc = ckpt['model'].nc# In[ ]:idx = 0
for k, v in model.state_dict().items():if "model.{}.".format(idx) in k:if idx < 29:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif idx < 42:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv2.".format(idx) in k:kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv3.".format(idx) in k:kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.dfl.".format(idx) in k:kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")else:while True:idx += 1if "model.{}.".format(idx) in k:breakif idx < 29:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif idx < 42:kr = k.replace("model.{}.".format(idx), "model.{}.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv2.".format(idx) in k:kr = k.replace("model.{}.cv2.".format(idx), "model.{}.cv4.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.cv3.".format(idx) in k:kr = k.replace("model.{}.cv3.".format(idx), "model.{}.cv5.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")elif "model.{}.dfl.".format(idx) in k:kr = k.replace("model.{}.dfl.".format(idx), "model.{}.dfl2.".format(idx+7))model.state_dict()[k] -= model.state_dict()[k]model.state_dict()[k] += ckpt['model'].state_dict()[kr]print(k, "perfectly matched!!")
_ = model.eval()# In[ ]:m_ckpt = {'model': model.half(),'optimizer': None,'best_fitness': None,'ema': None,'updates': None,'opt': None,'git': None,'date': None,'epoch': -1}
torch.save(m_ckpt, "./cityscapes-yolov9-e-converted.pt")