Bob's Blog

前些时候在使用yolov5训练出模型但遇到了加载时提示no module named models的问题，当时用取巧的方式绕过了这个问题，让训练出的模型能在自己的项目中被成功加载，但当时的解决方式只是个临时的，以后当目录结构有变化时容易导致继续修改，于是看了yolov5的代码和pytorch的官方链接。

https://www.byincd.com/bobjiang/article-01214/

https://pytorch.org/tutorials/beginner/saving_loading_models.html

pytorch官方推荐的是使用state dict。

除了在训练完成时直接保存为state dict，还可以加载用save保存的模型再保存为state dict。

import torch

m = torch.load("path/yolov5.pt")
torch.save(m["model"].state_dict(), "path/state.pth")

按照yolov5中train.py的参照可以这样加载，还好不用额外再写一份model class。但是必须用到自己项目的参数和配置定义，如下只是个例子。实际修改会比这个多，并不通用。

import torch
from models.yolo import Model

model = Model("path/ultralytics_yolov5/models/yolov5s.yaml", 3, 12, None).to("cpu")
model.load_state_dict(torch.load("path/state.pth"))
model.eval()

在attempt_load中的：

ckpt = torch.load(attempt_download(w), map_location=map_location)
model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())

如果此时加载的是state dict的文件内容，需要换成或加上条件分支：

ckpt = Model("path/yolov5s.yaml", 3, 12, None).to("cpu")
ckpt.load_state_dict(torch.load(weights, map_location=map_location))
ckpt.eval()
model.append(ckpt)

如果用的是自定义的训练配置，则需要指定到对应的yaml文件或者dict内容, 否则会得到非常多的size mismatch或missing key in state_dict的错误。

这里还有一个坑，是因为state dict未加上训练类别名，因此需要额外指定一下，否则能识别但无法给出具体的类别

ckpt.names = ['xx', 'xxx'...]

此时再度加载时就能跟使用pt文件一样的效果。解决了目录层次绑定的问题。