383 lines
14 KiB
Python
383 lines
14 KiB
Python
|
"""
|
|||
|
|
统一模型管理工具
|
|||
|
|
处理模型文件的统一命名、存储和检索
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import json
|
|||
|
|
import torch
|
|||
|
|
import glob
|
|||
|
|
from datetime import datetime
|
|||
|
|
from typing import List, Dict, Optional, Tuple
|
|||
|
|
from core.config import DEFAULT_MODEL_DIR
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ModelManager:
|
|||
|
|
"""统一模型管理器"""
|
|||
|
|
|
|||
|
|
def __init__(self, model_dir: str = DEFAULT_MODEL_DIR):
|
|||
|
|
self.model_dir = model_dir
|
|||
|
|
self.ensure_model_dir()
|
|||
|
|
|
|||
|
|
def ensure_model_dir(self):
|
|||
|
|
"""确保模型目录存在"""
|
|||
|
|
if not os.path.exists(self.model_dir):
|
|||
|
|
os.makedirs(self.model_dir)
|
|||
|
|
|
|||
|
|
def generate_model_filename(self,
|
|||
|
|
product_id: str,
|
|||
|
|
model_type: str,
|
|||
|
|
version: str,
|
|||
|
|
store_id: Optional[str] = None,
|
|||
|
|
training_mode: str = 'product',
|
|||
|
|
aggregation_method: Optional[str] = None) -> str:
|
|||
|
|
"""
|
|||
|
|
生成统一的模型文件名
|
|||
|
|
|
|||
|
|
格式规范:
|
|||
|
|
- 产品模式: {model_type}_product_{product_id}_{version}.pth
|
|||
|
|
- 店铺模式: {model_type}_store_{store_id}_{product_id}_{version}.pth
|
|||
|
|
- 全局模式: {model_type}_global_{product_id}_{aggregation_method}_{version}.pth
|
|||
|
|
"""
|
|||
|
|
if training_mode == 'store' and store_id:
|
|||
|
|
return f"{model_type}_store_{store_id}_{product_id}_{version}.pth"
|
|||
|
|
elif training_mode == 'global' and aggregation_method:
|
|||
|
|
return f"{model_type}_global_{product_id}_{aggregation_method}_{version}.pth"
|
|||
|
|
else:
|
|||
|
|
# 默认产品模式
|
|||
|
|
return f"{model_type}_product_{product_id}_{version}.pth"
|
|||
|
|
|
|||
|
|
def save_model(self,
|
|||
|
|
model_data: dict,
|
|||
|
|
product_id: str,
|
|||
|
|
model_type: str,
|
|||
|
|
version: str,
|
|||
|
|
store_id: Optional[str] = None,
|
|||
|
|
training_mode: str = 'product',
|
|||
|
|
aggregation_method: Optional[str] = None,
|
|||
|
|
product_name: Optional[str] = None) -> str:
|
|||
|
|
"""
|
|||
|
|
保存模型到统一位置
|
|||
|
|
|
|||
|
|
参数:
|
|||
|
|
model_data: 包含模型状态和配置的字典
|
|||
|
|
product_id: 产品ID
|
|||
|
|
model_type: 模型类型
|
|||
|
|
version: 版本号
|
|||
|
|
store_id: 店铺ID (可选)
|
|||
|
|
training_mode: 训练模式
|
|||
|
|
aggregation_method: 聚合方法 (可选)
|
|||
|
|
product_name: 产品名称 (可选)
|
|||
|
|
|
|||
|
|
返回:
|
|||
|
|
模型文件路径
|
|||
|
|
"""
|
|||
|
|
filename = self.generate_model_filename(
|
|||
|
|
product_id, model_type, version, store_id, training_mode, aggregation_method
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 统一保存到根目录，避免复杂的子目录结构
|
|||
|
|
model_path = os.path.join(self.model_dir, filename)
|
|||
|
|
|
|||
|
|
# 增强模型数据，添加管理信息
|
|||
|
|
enhanced_model_data = model_data.copy()
|
|||
|
|
enhanced_model_data.update({
|
|||
|
|
'model_manager_info': {
|
|||
|
|
'product_id': product_id,
|
|||
|
|
'product_name': product_name or product_id,
|
|||
|
|
'model_type': model_type,
|
|||
|
|
'version': version,
|
|||
|
|
'store_id': store_id,
|
|||
|
|
'training_mode': training_mode,
|
|||
|
|
'aggregation_method': aggregation_method,
|
|||
|
|
'created_at': datetime.now().isoformat(),
|
|||
|
|
'filename': filename
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 保存模型
|
|||
|
|
torch.save(enhanced_model_data, model_path)
|
|||
|
|
|
|||
|
|
print(f"模型已保存: {model_path}")
|
|||
|
|
return model_path
|
|||
|
|
|
|||
|
|
def list_models(self,
|
|||
|
|
product_id: Optional[str] = None,
|
|||
|
|
model_type: Optional[str] = None,
|
|||
|
|
store_id: Optional[str] = None,
|
|||
|
|
training_mode: Optional[str] = None,
|
|||
|
|
page: Optional[int] = None,
|
|||
|
|
page_size: Optional[int] = None) -> Dict:
|
|||
|
|
"""
|
|||
|
|
列出所有模型文件
|
|||
|
|
|
|||
|
|
参数:
|
|||
|
|
product_id: 产品ID过滤 (可选)
|
|||
|
|
model_type: 模型类型过滤 (可选)
|
|||
|
|
store_id: 店铺ID过滤 (可选)
|
|||
|
|
training_mode: 训练模式过滤 (可选)
|
|||
|
|
page: 页码，从1开始 (可选)
|
|||
|
|
page_size: 每页数量 (可选)
|
|||
|
|
|
|||
|
|
返回:
|
|||
|
|
包含模型列表和分页信息的字典
|
|||
|
|
"""
|
|||
|
|
models = []
|
|||
|
|
|
|||
|
|
# 搜索所有.pth文件
|
|||
|
|
pattern = os.path.join(self.model_dir, "*.pth")
|
|||
|
|
model_files = glob.glob(pattern)
|
|||
|
|
|
|||
|
|
for model_file in model_files:
|
|||
|
|
try:
|
|||
|
|
# 解析文件名
|
|||
|
|
filename = os.path.basename(model_file)
|
|||
|
|
model_info = self.parse_model_filename(filename)
|
|||
|
|
|
|||
|
|
if not model_info:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 尝试从模型文件中读取额外信息
|
|||
|
|
try:
|
|||
|
|
# Try with weights_only=False first for backward compatibility
|
|||
|
|
try:
|
|||
|
|
model_data = torch.load(model_file, map_location='cpu', weights_only=False)
|
|||
|
|
except Exception:
|
|||
|
|
# If that fails, try with weights_only=True (newer PyTorch versions)
|
|||
|
|
model_data = torch.load(model_file, map_location='cpu', weights_only=True)
|
|||
|
|
|
|||
|
|
if 'model_manager_info' in model_data:
|
|||
|
|
# 使用新的管理信息
|
|||
|
|
manager_info = model_data['model_manager_info']
|
|||
|
|
model_info.update(manager_info)
|
|||
|
|
|
|||
|
|
# 添加评估指标
|
|||
|
|
if 'metrics' in model_data:
|
|||
|
|
model_info['metrics'] = model_data['metrics']
|
|||
|
|
|
|||
|
|
# 添加配置信息
|
|||
|
|
if 'config' in model_data:
|
|||
|
|
model_info['config'] = model_data['config']
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"读取模型文件失败 {model_file}: {e}")
|
|||
|
|
# Continue with just the filename-based info
|
|||
|
|
|
|||
|
|
# 应用过滤器
|
|||
|
|
if product_id and model_info.get('product_id') != product_id:
|
|||
|
|
continue
|
|||
|
|
if model_type and model_info.get('model_type') != model_type:
|
|||
|
|
continue
|
|||
|
|
if store_id and model_info.get('store_id') != store_id:
|
|||
|
|
continue
|
|||
|
|
if training_mode and model_info.get('training_mode') != training_mode:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 添加文件信息
|
|||
|
|
model_info['filename'] = filename
|
|||
|
|
model_info['file_path'] = model_file
|
|||
|
|
model_info['file_size'] = os.path.getsize(model_file)
|
|||
|
|
model_info['modified_at'] = datetime.fromtimestamp(
|
|||
|
|
os.path.getmtime(model_file)
|
|||
|
|
).isoformat()
|
|||
|
|
|
|||
|
|
models.append(model_info)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"处理模型文件失败 {model_file}: {e}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 按创建时间排序（最新的在前）
|
|||
|
|
models.sort(key=lambda x: x.get('created_at', x.get('modified_at', '')), reverse=True)
|
|||
|
|
|
|||
|
|
# 计算分页信息
|
|||
|
|
total_count = len(models)
|
|||
|
|
|
|||
|
|
# 如果没有指定分页参数，返回所有数据
|
|||
|
|
if page is None or page_size is None:
|
|||
|
|
return {
|
|||
|
|
'models': models,
|
|||
|
|
'pagination': {
|
|||
|
|
'total': total_count,
|
|||
|
|
'page': 1,
|
|||
|
|
'page_size': total_count,
|
|||
|
|
'total_pages': 1,
|
|||
|
|
'has_next': False,
|
|||
|
|
'has_previous': False
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 应用分页
|
|||
|
|
total_pages = (total_count + page_size - 1) // page_size if page_size > 0 else 1
|
|||
|
|
start_index = (page - 1) * page_size
|
|||
|
|
end_index = start_index + page_size
|
|||
|
|
|
|||
|
|
paginated_models = models[start_index:end_index]
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'models': paginated_models,
|
|||
|
|
'pagination': {
|
|||
|
|
'total': total_count,
|
|||
|
|
'page': page,
|
|||
|
|
'page_size': page_size,
|
|||
|
|
'total_pages': total_pages,
|
|||
|
|
'has_next': page < total_pages,
|
|||
|
|
'has_previous': page > 1
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def parse_model_filename(self, filename: str) -> Optional[Dict]:
|
|||
|
|
"""
|
|||
|
|
解析模型文件名，提取模型信息
|
|||
|
|
|
|||
|
|
支持的格式:
|
|||
|
|
- {model_type}_product_{product_id}_{version}.pth
|
|||
|
|
- {model_type}_store_{store_id}_{product_id}_{version}.pth
|
|||
|
|
- {model_type}_global_{product_id}_{aggregation_method}_{version}.pth
|
|||
|
|
- 旧格式兼容
|
|||
|
|
"""
|
|||
|
|
if not filename.endswith('.pth'):
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
base_name = filename.replace('.pth', '')
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 新格式解析
|
|||
|
|
if '_product_' in base_name:
|
|||
|
|
# 产品模式: model_type_product_product_id_version
|
|||
|
|
parts = base_name.split('_product_')
|
|||
|
|
model_type = parts[0]
|
|||
|
|
rest = parts[1]
|
|||
|
|
|
|||
|
|
# 分离产品ID和版本
|
|||
|
|
if '_v' in rest:
|
|||
|
|
last_v_index = rest.rfind('_v')
|
|||
|
|
product_id = rest[:last_v_index]
|
|||
|
|
version = rest[last_v_index+1:]
|
|||
|
|
else:
|
|||
|
|
product_id = rest
|
|||
|
|
version = 'v1'
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'model_type': model_type,
|
|||
|
|
'product_id': product_id,
|
|||
|
|
'version': version,
|
|||
|
|
'training_mode': 'product',
|
|||
|
|
'store_id': None,
|
|||
|
|
'aggregation_method': None
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
elif '_store_' in base_name:
|
|||
|
|
# 店铺模式: model_type_store_store_id_product_id_version
|
|||
|
|
parts = base_name.split('_store_')
|
|||
|
|
model_type = parts[0]
|
|||
|
|
rest = parts[1]
|
|||
|
|
|
|||
|
|
# 分离店铺ID、产品ID和版本
|
|||
|
|
rest_parts = rest.split('_')
|
|||
|
|
if len(rest_parts) >= 3:
|
|||
|
|
store_id = rest_parts[0]
|
|||
|
|
if rest_parts[-1].startswith('v'):
|
|||
|
|
# 最后一部分是版本号
|
|||
|
|
version = rest_parts[-1]
|
|||
|
|
product_id = '_'.join(rest_parts[1:-1])
|
|||
|
|
else:
|
|||
|
|
version = 'v1'
|
|||
|
|
product_id = '_'.join(rest_parts[1:])
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'model_type': model_type,
|
|||
|
|
'product_id': product_id,
|
|||
|
|
'version': version,
|
|||
|
|
'training_mode': 'store',
|
|||
|
|
'store_id': store_id,
|
|||
|
|
'aggregation_method': None
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
elif '_global_' in base_name:
|
|||
|
|
# 全局模式: model_type_global_product_id_aggregation_method_version
|
|||
|
|
parts = base_name.split('_global_')
|
|||
|
|
model_type = parts[0]
|
|||
|
|
rest = parts[1]
|
|||
|
|
|
|||
|
|
rest_parts = rest.split('_')
|
|||
|
|
if len(rest_parts) >= 3:
|
|||
|
|
if rest_parts[-1].startswith('v'):
|
|||
|
|
# 最后一部分是版本号
|
|||
|
|
version = rest_parts[-1]
|
|||
|
|
aggregation_method = rest_parts[-2]
|
|||
|
|
product_id = '_'.join(rest_parts[:-2])
|
|||
|
|
else:
|
|||
|
|
version = 'v1'
|
|||
|
|
aggregation_method = rest_parts[-1]
|
|||
|
|
product_id = '_'.join(rest_parts[:-1])
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'model_type': model_type,
|
|||
|
|
'product_id': product_id,
|
|||
|
|
'version': version,
|
|||
|
|
'training_mode': 'global',
|
|||
|
|
'store_id': None,
|
|||
|
|
'aggregation_method': aggregation_method
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 兼容旧格式
|
|||
|
|
else:
|
|||
|
|
# 尝试解析其他格式
|
|||
|
|
if 'model_product_' in base_name:
|
|||
|
|
parts = base_name.split('_model_product_')
|
|||
|
|
model_type = parts[0]
|
|||
|
|
product_part = parts[1]
|
|||
|
|
|
|||
|
|
if '_v' in product_part:
|
|||
|
|
last_v_index = product_part.rfind('_v')
|
|||
|
|
product_id = product_part[:last_v_index]
|
|||
|
|
version = product_part[last_v_index+1:]
|
|||
|
|
else:
|
|||
|
|
product_id = product_part
|
|||
|
|
version = 'v1'
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'model_type': model_type,
|
|||
|
|
'product_id': product_id,
|
|||
|
|
'version': version,
|
|||
|
|
'training_mode': 'product',
|
|||
|
|
'store_id': None,
|
|||
|
|
'aggregation_method': None
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"解析文件名失败 {filename}: {e}")
|
|||
|
|
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def delete_model(self, model_file: str) -> bool:
|
|||
|
|
"""删除模型文件"""
|
|||
|
|
try:
|
|||
|
|
if os.path.exists(model_file):
|
|||
|
|
os.remove(model_file)
|
|||
|
|
print(f"已删除模型文件: {model_file}")
|
|||
|
|
return True
|
|||
|
|
else:
|
|||
|
|
print(f"模型文件不存在: {model_file}")
|
|||
|
|
return False
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"删除模型文件失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def get_model_by_id(self, model_id: str) -> Optional[Dict]:
|
|||
|
|
"""根据模型ID获取模型信息"""
|
|||
|
|
models = self.list_models()
|
|||
|
|
for model in models:
|
|||
|
|
if model.get('filename', '').replace('.pth', '') == model_id:
|
|||
|
|
return model
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局模型管理器实例
|
|||
|
|
# 确保使用项目根目录的saved_models，而不是相对于当前工作目录
|
|||
|
|
import os
|
|||
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|||
|
|
project_root = os.path.dirname(os.path.dirname(current_dir)) # 向上两级到项目根目录
|
|||
|
|
absolute_model_dir = os.path.join(project_root, 'saved_models')
|
|||
|
|
model_manager = ModelManager(absolute_model_dir)
|