# 如何向系统添加新模型 本指南详细说明了如何向本预测系统添加一个全新的预测模型。系统采用灵活的插件式架构，集成新模型的过程非常模块化，主要围绕 **模型（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` 进行训练和预测时，系统将自动调用您刚刚编写和注册的所有相应逻辑。