如何加快自定义LSTM DQN金融时序强化学习模型的训练速度

核心问题排查（GPU无提速原因） #

• 你的LSTM输入格式不符合cuDNN加速要求：Keras默认LSTM的输入格式为(batch, 时间步长, 特征数)，你当前设置的(no_of_features, lookback_window)会跳过cuDNN加速，实际运行的是CPU版本的LSTM算子，因此切换GPU没有效果。
• cuDNN对LSTM的加速要求激活函数为默认的tanh、循环激活为sigmoid，你在LSTM层后额外添加的LeakyReLU层也会导致无法调用加速算子。

可直接落地的优化点 #

1. 模型层修改适配GPU加速 #

def create_DL_model():
  model = Sequential()
  # 调整输入维度为(时间步长, 特征数)，使用默认激活触发cuDNN加速
  model.add(LSTM(64,input_shape=(lookback_window, no_of_features), return_sequences=True))
  model.add(LSTM(32))
  model.add(Dense(no_of_actions))
  model.compile(loss='mse',optimizer = Adam(learning_rate = learning_rate))
  return model

注意：需要同步调整环境输出的state格式，把原来的(9,60)转置为(60,9)

2. 替换高开销的API调用 #

• 单样本推理时不要用predict，换成直接调用模型的方式，减少调度开销：

# 原来的写法
action = np.argmax(main_net.predict(state)[0])
# 替换为
action = np.argmax(main_net(state, training=False).numpy()[0])

• 单批次训练时不要用fit，换成train_on_batch，省去回调、日志等额外开销：

# 原来的写法
main_net.fit(np.array(X),np.array(y),batch_size = batch_size,verbose=0)
# 替换为
main_net.train_on_batch(np.array(X), np.array(y))

3. 减少训练和推理的调用频率 #

• 不要每个时间步都执行训练，改为每N步训练一次，比如设置每4步训练一次，训练开销直接降低75%，对DQN效果影响极小。
• 不用每个时间步都做单独的Q值推理，可攒多个状态批量推理，进一步降低GPU调度开销。

4. 移除Q值计算的Python循环 #

原来的batch内循环计算Q值可以改为全向量化操作，大幅提升速度：

# 采样后直接批量计算，不需要for循环
done_mask = np.array([1 if not trans[4] else 0 for trans in batch])
rewards = np.array([trans[2] for trans in batch])
actions = np.array([trans[1] for trans in batch])

max_future_q = np.max(future_Q_list, axis=1)
new_q = rewards + discount * max_future_q * done_mask

Q_list[np.arange(batch_size), actions] = new_q
X = np.array([trans[0][0] for trans in batch])
y = Q_list

5. 其他小优化 #

• 修正代码笔误：target_net.set_weights(tmain_net.get_weights())中的tmain_net应为main_net。
• 不需要每次训练都加with tensorflow.device('/gpu:0')，只要TensorFlow检测到GPU会默认调用，重复指定反而增加开销。

做完以上优化后，训练速度至少能提升10倍以上，2个episode的运行时间会降到1分钟以内。

内容的提问来源于stack exchange，提问作者Ryan aA