在做工业项目时,yolov5是比较好的检测算法,
但是如果一个项目中需要多个版本的yolov5检测项目,检测不同的物体,
加载不同版本的网络会报错
self.weights='yolov5s.pth'
model = torch.load(self.weights)
stride = int(model.stride.max()) # model stride
原因有:
1.网络结构不一样,比如激活函数不一样,层内容不一样
2.检测层Detect不一样。
解决方法:jit.trace方式。
导出:
self.model.model[-1].export = True
self.model.eval()
self.model=self.model.float()
x = torch.randn(1, 3, 544, 960).float().cuda()
traced_cell = torch.jit.trace(self.model, (x))
torch.jit.save(traced_cell, "jit_face_detect.pth"