如何用Python构建联邦学习LSTM模型实现股票收盘价预测

你之前定义的create_tf_dataset_for_client_fn需要实现时间序列样本构造逻辑，用历史N个时间步的特征预测下一个时间步的Close收盘价，参考实现如下：

import tensorflow as tf

# 超参数设置
TIME_STEPS = 30  # 用过去30天的数据预测下一天收盘价
FEATURE_COLS = ['Open', 'High', 'Close', 'Volume']
TARGET_COL = 'Close'
BATCH_SIZE = 32

def create_tf_dataset_for_client_fn(client_id):
    # 取出对应客户组的原始数据，按日期排序保证时间序列顺序
    client_df = df[df[client_id_colname] == client_id].sort_values('Date').reset_index(drop=True)
    # 提取特征和目标值
    features = client_df[FEATURE_COLS].values
    targets = client_df[TARGET_COL].values
    
    # 滑窗构造样本
    X, y = [], []
    for i in range(TIME_STEPS, len(client_df)):
        X.append(features[i-TIME_STEPS:i])
        y.append(targets[i])
    
    # 转为TF数据集，按时间顺序切分训练/测试集，禁止打乱
    dataset = tf.data.Dataset.from_tensor_slices((X, y))
    train_size = int(len(dataset) * 0.8)
    # 训练场景取前80%，测试场景取后20%
    if 'train' in tf.data.Dataset.context:
        dataset = dataset.take(train_size)
    else:
        dataset = dataset.skip(train_size)
    
    # 每个客户端独立做特征归一化，避免跨客户端数据泄露
    def normalize(x, y):
        x = tf.cast(x, tf.float32) / tf.reduce_max(x, axis=0)
        y = tf.cast(y, tf.float32) / tf.reduce_max(y)
        return x, y
    
    return dataset.map(normalize).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

定义符合TFF要求的Keras LSTM模型，输入维度为(时间步长, 特征数)，输出为单值的收盘价预测结果：

def build_lstm_model():
    model = tf.keras.Sequential([
        tf.keras.layers.LSTM(64, return_sequences=True, input_shape=(TIME_STEPS, len(FEATURE_COLS))),
        tf.keras.layers.LSTM(32, return_sequences=False),
        tf.keras.layers.Dense(16, activation='relu'),
        tf.keras.layers.Dense(1)
    ])
    model.compile(
        loss=tf.keras.losses.MeanSquaredError(),
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
        metrics=[tf.keras.metrics.MeanAbsoluteError(), tf.keras.metrics.RootMeanSquaredError()]
    )
    return model

# 验证模型输入输出格式是否符合要求
test_model = build_lstm_model()
sample_input = tf.random.normal((1, TIME_STEPS, len(FEATURE_COLS)))
print(test_model(sample_input).shape) # 输出应为(1,1)

将Keras模型转为TFF可识别的格式，使用联邦平均算法完成训练流程定义：

import tensorflow_federated as tff

# 定义TFF模型包装器
def model_fn():
    keras_model = build_lstm_model()
    loss = tf.keras.losses.MeanSquaredError()
    metrics = [tf.keras.metrics.MeanAbsoluteError(), tf.keras.metrics.RootMeanSquaredError()]
    return tff.learning.models.from_keras_model(
        keras_model,
        input_spec=train_data.create_tf_dataset_for_client(train_data.client_ids[0]).element_spec,
        loss=loss,
        metrics=metrics
    )

# 构建加权联邦平均训练过程
training_process = tff.learning.algorithms.build_weighted_fed_avg(
    model_fn=model_fn,
    client_optimizer_fn=lambda: tf.keras.optimizers.Adam(learning_rate=0.001),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0)
)

# 构建联邦评估过程
eval_process = tff.learning.build_federated_evaluation(model_fn=model_fn)

初始化训练状态，迭代多轮联邦训练，每轮训练后可在测试集上评估效果：

# 初始化服务端模型状态
state = training_process.initialize()

# 训练轮次超参数
NUM_ROUNDS = 20

for round_num in range(1, NUM_ROUNDS+1):
    # 每轮选中所有客户端参与训练，也可根据需求按比例采样
    train_client_datasets = [train_data.create_tf_dataset_for_client(cid) for cid in train_data.client_ids]
    # 执行一轮训练
    result = training_process.next(state, train_client_datasets)
    state = result.state
    train_metrics = result.metrics
    print(f'第{round_num}轮训练结束，训练MSE：{train_metrics["client_work"]["train"]["loss"]:.4f}，训练MAE：{train_metrics["client_work"]["train"]["mean_absolute_error"]:.4f}')
    
    # 每5轮做一次测试集评估
    if round_num %5 ==0:
        test_client_datasets = [test_data.create_tf_dataset_for_client(cid) for cid in test_data.client_ids]
        eval_metrics = eval_process(state.model, test_client_datasets)
        print(f'第{round_num}轮测试结果，测试MSE：{eval_metrics["eval"]["loss"]:.4f}，测试MAE：{eval_metrics["eval"]["mean_absolute_error"]:.4f}')

• 时间序列数据切分训练测试集必须严格按时间顺序划分，禁止随机打乱，避免未来数据泄露
• 特征归一化的统计值只能用对应客户端的训练集数据计算，不能用到测试集或其他客户端的数据
• 可以根据实际效果调整LSTM层数、隐藏单元数、时间步长、学习率等超参数

内容的提问来源于stack exchange，提问作者Lim Shi Hao