如何向系统添加新模型

本指南详细说明了如何向本预测系统添加一个全新的预测模型。系统采用灵活的插件式架构，集成新模型的过程非常模块化，主要围绕 模型（Model）、训练器（Trainer） 和 预测器（Predictor） 这三个核心组件进行。

核心理念

系统的核心是 models/model_registry.py，它维护了两个独立的注册表：一个用于训练函数，另一个用于预测函数。添加新模型的本质就是：

定义模型：创建模型的架构。
创建训练器：编写一个函数来训练这个模型，并将其注册到训练器注册表。
集成预测器：确保系统知道如何加载模型并用它来预测，然后将预测逻辑注册到预测器注册表。

第 1 步：定义模型架构

首先，您需要在 ShopTRAINING/server/models/ 目录下创建一个新的 Python 文件来定义您的模型。

示例：创建 ShopTRAINING/server/models/my_new_model.py

如果您的新模型是基于 PyTorch 的，它应该是一个继承自 torch.nn.Module 的类。

# 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 函数。

# 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

让系统知道如何构建您的模型实例

在 load_model_and_predict 函数中，有一个 if/elif 结构用于根据模型类型实例化不同的模型。您需要为 MyNewModel 添加一个新的分支。

# 在 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()

注册预测逻辑

如果您的模型是一个标准的 PyTorch 模型，并且其预测逻辑与现有的模型（如 Transformer, KAN）相同，您可以直接复用 default_pytorch_predictor。只需在文件末尾添加一行注册代码即可。
```
# 在 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 进行训练和预测时，系统将自动调用您刚刚编写和注册的所有相应逻辑。

8.3 KiB Raw Permalink Blame History Unescape Escape

如何向系统添加新模型

核心理念

第 1 步：定义模型架构

第 2 步：创建模型训练器

第 3 步：集成模型预测器

总结

8.3 KiB

Raw Permalink Blame History