将任意图像增强模型与ultralytics中任意模型进行结合,支持自定义图像增强模块的loss,实现端到端训练
图像增强模型算法版本较多,同时ultralytics库的迭代也比较大,直接将图像增强模型算法添加到ultralytics库中可能只是当前版本支持,在ultralytics库升级后,又需要重新配置。故而分享一种方式进行实现,支持任意图像增强模型与ultralytics库中任务模型(如yolov8,yolov9,yolov10等模型)进行结合。支持模型的训练、验证与导出。
其核心就是,对模型结构进行替换,同时确保调用ultralytics库接口的train函数时模型结构替换依旧有效。
虽然没有对gdip模块进行测试,但确认这种修改方式是支持实现gdip-yolov10的,也支持实现ia-yolov10。
1、保存图像增强模型
如sci++项目,将模型结构带权重保存为pth文件即可,具体参考示意如下所示
from model import Finetunemodel
import torchif __name__ ==<