pytorch保存、加载模型的方法

一、保存模型的 结构+权重

保存模型

torch.save(model, 'model.pt')

加载模型

model = torch.load('model.pt')

 二、只保存模型的 权重

 保存模型

torch.save(model.state_dict(), 'weights.pt')

加载模型

model = MyModel()  # 需要先定义模型结构
model.load_state_dict(torch.load('weights.pt'))  # 再把权重加载到模型结构上

YOLO模型结构和权重的关系

发现很多朋友把模型文件和权重文件搞混了，在此基于YOLOv5/v8形象的解释一下：

1、yolov8s.yaml，里面存了模型结构，可以认为是骨架

2、yolov8s.pt，这是权重文件，是训练得到的，里面保存了每一层训练出来的参数，比如卷积核

总结一下就是：

从零训练：用yaml构建模型      -----训练----->       得到pt

用预训练权重训练：加载pt      -----训练----->       得到pt

改模型结构加载预训练权重训练：用yaml构建模型, 加载pt权重      -----训练----->       得到pt

注：

改动过的模型，建议不要加载预训练权重

因为yolov8s.pt的权重匹配的是yolov8s.yaml的结构，你改动了模型，结构变了，大部分层的结构和权重都是对不上的，相当于你把170的裤子硬穿到180的人身上，天安门可能就对不上了