56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import sys
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
def attempt_load(weights, device=''):
|
|||
|
|
"""尝试加载YOLOv5模型"""
|
|||
|
|
# 加载模型
|
|||
|
|
model = torch.load(weights, map_location=device)
|
|||
|
|
|
|||
|
|
# 确定模型格式
|
|||
|
|
if isinstance(model, dict):
|
|||
|
|
if 'model' in model: # state_dict格式
|
|||
|
|
model = model['model']
|
|||
|
|
elif 'state_dict' in model: # state_dict格式
|
|||
|
|
model = model['state_dict']
|
|||
|
|
|
|||
|
|
# 如果是state_dict，则需要创建模型架构
|
|||
|
|
if isinstance(model, dict):
|
|||
|
|
print("警告：加载的是权重字典，尝试创建默认模型结构")
|
|||
|
|
from models.yolov5_model import YOLOv5
|
|||
|
|
model_arch = YOLOv5()
|
|||
|
|
model_arch.load_state_dict(model)
|
|||
|
|
model = model_arch
|
|||
|
|
|
|||
|
|
# 设置为评估模式
|
|||
|
|
if isinstance(model, nn.Module):
|
|||
|
|
model.eval()
|
|||
|
|
|
|||
|
|
# 检查是否有类别信息
|
|||
|
|
if not hasattr(model, 'names') or not model.names:
|
|||
|
|
print("模型没有类别信息，尝试加载默认类别")
|
|||
|
|
# 设置通用类别
|
|||
|
|
model.names = ['object']
|
|||
|
|
|
|||
|
|
return model
|
|||
|
|
|
|||
|
|
class YOLOv5:
|
|||
|
|
"""简化版YOLOv5模型结构，用于加载权重"""
|
|||
|
|
def __init__(self):
|
|||
|
|
super(YOLOv5, self).__init__()
|
|||
|
|
self.names = [] # 类别名称
|
|||
|
|
# 这里应该添加真实的网络结构
|
|||
|
|
# 但为了简单起见，我们只提供一个占位符
|
|||
|
|
# 在实际使用中，您应该实现完整的网络架构
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
# 这里应该是实际的前向传播逻辑
|
|||
|
|
# 这只是一个占位符
|
|||
|
|
raise NotImplementedError("这是一个占位符模型，请使用完整的YOLOv5模型实现")
|
|||
|
|
|
|||
|
|
def load_state_dict(self, state_dict):
|
|||
|
|
print("尝试加载模型权重")
|
|||
|
|
# 实际的权重加载逻辑
|
|||
|
|
# 这只是一个占位符
|
|||
|
|
return self
|