222 lines
8.3 KiB
Markdown
222 lines
8.3 KiB
Markdown
# 如何向系统添加新模型
|
||
|
||
本指南详细说明了如何向本预测系统添加一个全新的预测模型。系统采用灵活的插件式架构，集成新模型的过程非常模块化，主要围绕 **模型（Model）**、**训练器（Trainer）** 和 **预测器（Predictor）** 这三个核心组件进行。
|
||
|
||
## 核心理念
|
||
|
||
系统的核心是 `models/model_registry.py`，它维护了两个独立的注册表：一个用于训练函数，另一个用于预测函数。添加新模型的本质就是：
|
||
|
||
1. **定义模型**：创建模型的架构。
|
||
2. **创建训练器**：编写一个函数来训练这个模型，并将其注册到训练器注册表。
|
||
3. **集成预测器**：确保系统知道如何加载模型并用它来预测，然后将预测逻辑注册到预测器注册表。
|
||
|
||
---
|
||
|
||
## 第 1 步：定义模型架构
|
||
|
||
首先，您需要在 `ShopTRAINING/server/models/` 目录下创建一个新的 Python 文件来定义您的模型。
|
||
|
||
**示例：创建 `ShopTRAINING/server/models/my_new_model.py`**
|
||
|
||
如果您的新模型是基于 PyTorch 的，它应该是一个继承自 `torch.nn.Module` 的类。
|
||
|
||
```python
|
||
# file: ShopTRAINING/server/models/my_new_model.py
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
class MyNewModel(nn.Module):
|
||
def __init__(self, input_features, hidden_size, output_sequence_length):
|
||
"""
|
||
定义模型的层和结构。
|
||
"""
|
||
super(MyNewModel, self).__init__()
|
||
self.layer1 = nn.Linear(input_features, hidden_size)
|
||
self.relu = nn.ReLU()
|
||
self.layer2 = nn.Linear(hidden_size, output_sequence_length)
|
||
# ... 可添加更复杂的结构
|
||
|
||
def forward(self, x):
|
||
"""
|
||
定义数据通过模型的前向传播路径。
|
||
x 的形状通常是 (batch_size, sequence_length, num_features)
|
||
"""
|
||
# 确保输入是正确的形状
|
||
# 例如，对于简单的线性层，可能需要展平
|
||
batch_size, seq_len, features = x.shape
|
||
x = x.view(batch_size * seq_len, features) # 展平
|
||
|
||
out = self.layer1(x)
|
||
out = self.relu(out)
|
||
out = self.layer2(out)
|
||
|
||
# 恢复形状以匹配输出
|
||
out = out.view(batch_size, seq_len, -1)
|
||
# 通常我们只关心序列的最后一个预测
|
||
return out[:, -1, :]
|
||
```
|
||
|
||
---
|
||
|
||
## 第 2 步：创建模型训练器
|
||
|
||
接下来，在 `ShopTRAINING/server/trainers/` 目录下创建一个新的训练器文件。这个文件负责模型的整个训练、评估和保存流程。
|
||
|
||
**示例：创建 `ShopTRAINING/server/trainers/my_new_model_trainer.py`**
|
||
|
||
这个训练函数需要遵循系统中其他训练器（如 `xgboost_trainer.py`）的统一函数签名，并使用 `@register_trainer` 装饰器或在文件末尾调用 `register_trainer` 函数。
|
||
|
||
```python
|
||
# file: ShopTRAINING/server/trainers/my_new_model_trainer.py
|
||
|
||
import torch
|
||
import torch.optim as optim
|
||
from models.model_registry import register_trainer
|
||
from utils.model_manager import model_manager
|
||
from analysis.metrics import evaluate_model
|
||
from models.my_new_model import MyNewModel # 导入您的新模型
|
||
|
||
# 遵循系统的标准函数签名
|
||
def train_with_mynewmodel(product_id, model_identifier, product_df, store_id, training_mode, aggregation_method, epochs, sequence_length, forecast_horizon, model_dir, **kwargs):
|
||
print(f"🚀 MyNewModel 训练器启动: model_identifier='{model_identifier}'")
|
||
|
||
# --- 1. 数据准备 ---
|
||
# (此处省略了数据加载、标准化和创建数据集的详细代码,
|
||
# 您可以参考 xgboost_trainer.py 或其他训练器中的实现)
|
||
# ...
|
||
# 假设您已准备好 trainX, trainY, testX, testY, scaler_y 等变量
|
||
# trainX = ...
|
||
# trainY = ...
|
||
# testX = ...
|
||
# testY = ...
|
||
# scaler_y = ...
|
||
# features = [...]
|
||
|
||
# --- 2. 实例化模型和优化器 ---
|
||
input_dim = trainX.shape[2] # 获取特征数量
|
||
hidden_size = 64 # 示例超参数
|
||
|
||
model = MyNewModel(
|
||
input_features=input_dim,
|
||
hidden_size=hidden_size,
|
||
output_sequence_length=forecast_horizon
|
||
)
|
||
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
||
criterion = torch.nn.MSELoss()
|
||
|
||
# --- 3. 训练循环 ---
|
||
print("开始训练 MyNewModel...")
|
||
for epoch in range(epochs):
|
||
model.train()
|
||
optimizer.zero_grad()
|
||
outputs = model(trainX)
|
||
loss = criterion(outputs, trainY)
|
||
loss.backward()
|
||
optimizer.step()
|
||
if (epoch + 1) % 10 == 0:
|
||
print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')
|
||
|
||
# --- 4. 模型评估 ---
|
||
model.eval()
|
||
with torch.no_grad():
|
||
test_pred_scaled = model(testX)
|
||
|
||
# 反标准化并计算指标
|
||
# ... (参考其他训练器)
|
||
metrics = {'rmse': 0.0, 'mae': 0.0, 'r2': 0.0, 'mape': 0.0} # 示例
|
||
|
||
# --- 5. 模型保存 ---
|
||
model_data = {
|
||
'model_state_dict': model.state_dict(),
|
||
'scaler_X': None, # 替换为您的 scaler_X
|
||
'scaler_y': scaler_y,
|
||
'config': {
|
||
'model_type': 'mynewmodel', # **关键**: 使用唯一的模型名称
|
||
'input_dim': input_dim,
|
||
'hidden_size': hidden_size,
|
||
'sequence_length': sequence_length,
|
||
'forecast_horizon': forecast_horizon,
|
||
'features': features
|
||
},
|
||
'metrics': metrics
|
||
}
|
||
|
||
model_manager.save_model(
|
||
model_data=model_data,
|
||
product_id=product_id,
|
||
model_type='mynewmodel', # **关键**: 再次确认模型名称
|
||
# ... 其他参数
|
||
)
|
||
|
||
print("✅ MyNewModel 模型训练并保存完成！")
|
||
return model, metrics, "v1", "path/to/model" # 返回值遵循统一格式
|
||
|
||
# --- 关键步骤: 将训练器注册到系统中 ---
|
||
register_trainer('mynewmodel', train_with_mynewmodel)
|
||
```
|
||
|
||
---
|
||
|
||
## 第 3 步：集成模型预测器
|
||
|
||
最后，您需要让系统知道如何加载和使用您的新模型进行预测。这需要在 `ShopTRAINING/server/predictors/model_predictor.py` 中进行两处修改。
|
||
|
||
**文件: `ShopTRAINING/server/predictors/model_predictor.py`**
|
||
|
||
1. **让系统知道如何构建您的模型实例**
|
||
|
||
在 `load_model_and_predict` 函数中，有一个 `if/elif` 结构用于根据模型类型实例化不同的模型。您需要为 `MyNewModel` 添加一个新的分支。
|
||
|
||
```python
|
||
# 在 model_predictor.py 中
|
||
|
||
# 首先，导入您的新模型类
|
||
from models.my_new_model import MyNewModel
|
||
|
||
# ... 在 load_model_and_predict 函数内部 ...
|
||
|
||
# ... 其他模型的 elif 分支 ...
|
||
elif loaded_model_type == 'tcn':
|
||
model = TCNForecaster(...)
|
||
|
||
# vvv 添加这个新的分支 vvv
|
||
elif loaded_model_type == 'mynewmodel':
|
||
model = MyNewModel(
|
||
input_features=config['input_dim'],
|
||
hidden_size=config['hidden_size'],
|
||
output_sequence_length=config['forecast_horizon']
|
||
).to(DEVICE)
|
||
# ^^^ 添加结束 ^^^
|
||
|
||
else:
|
||
raise ValueError(f"不支持的模型类型: {loaded_model_type}")
|
||
|
||
model.load_state_dict(checkpoint['model_state_dict'])
|
||
model.eval()
|
||
```
|
||
|
||
2. **注册预测逻辑**
|
||
|
||
如果您的模型是一个标准的 PyTorch 模型，并且其预测逻辑与现有的模型（如 Transformer, KAN）相同，您可以直接复用 `default_pytorch_predictor`。只需在文件末尾添加一行注册代码即可。
|
||
|
||
```python
|
||
# 在 model_predictor.py 文件末尾
|
||
|
||
# ...
|
||
# 将增强后的默认预测器也注册给xgboost
|
||
register_predictor('xgboost', default_pytorch_predictor)
|
||
|
||
# vvv 添加这行代码 vvv
|
||
# 让 'mynewmodel' 也使用通用的 PyTorch 预测器
|
||
register_predictor('mynewmodel', default_pytorch_predictor)
|
||
# ^^^ 添加结束 ^^^
|
||
```
|
||
|
||
如果您的模型需要特殊的预测逻辑（例如，像 XGBoost 那样有不同的输入格式或调用方式），您可以复制 `default_pytorch_predictor` 创建一个新函数，修改其内部逻辑，然后将新函数注册给 `'mynewmodel'`。
|
||
|
||
---
|
||
|
||
## 总结
|
||
|
||
完成以上三个步骤后，您的新模型 `MyNewModel` 就已完全集成到系统中了。系统会自动在 `trainers` 目录中发现您的新训练器。当您通过 API 或界面选择 `mynewmodel` 进行训练和预测时，系统将自动调用您刚刚编写和注册的所有相应逻辑。 |