为什么这个 tensorflow 训练需要这么长时间？

Question

我正在通过深度强化学习实战一书学习 DRL。在第 3 章中，他们展示了简单游戏 Gridworld（instructions here，在规则部分）以及 PyTorch 中的相应代码。

我对代码进行了试验，训练网络用时不到 3 分钟，获胜率为 89%（训练后赢得 100 场比赛中的 89 场）。

作为练习，我已将代码迁移到 tensorflow。所有代码都是here。

问题是，使用我的 tensorflow 端口，以 84% 的胜率训练网络需要将近 2 个小时。两个版本都使用唯一的 CPU 进行训练（我没有 GPU）

训练损失数字似乎是正确的，也是获胜率（我们必须考虑到游戏是随机的，可能有不可能的状态）。问题是整个流程的表现。

我做错了什么，但是什么？

主要区别在于训练循环，在火炬中是这样的：

        loss_fn = torch.nn.MSELoss()
        learning_rate = 1e-3
        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
        ....
        Q1 = model(state1_batch) 
        with torch.no_grad():
            Q2 = model2(state2_batch) #B
        
        Y = reward_batch + gamma * ((1-done_batch) * torch.max(Q2,dim=1)[0])
        X = Q1.gather(dim=1,index=action_batch.long().unsqueeze(dim=1)).squeeze()
        loss = loss_fn(X, Y.detach())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

在 tensorflow 版本中：

        loss_fn = tf.keras.losses.MSE
        learning_rate = 1e-3
        optimizer = tf.keras.optimizers.Adam(learning_rate)
        ...
        Q2 = model2(state2_batch) #B
        with tf.GradientTape() as tape:
            Q1 = model(state1_batch)
            Y = reward_batch + gamma * ((1-done_batch) * tf.math.reduce_max(Q2, axis=1))
            X = [Q1[i][action_batch[i]] for i in range(len(action_batch))]
            loss = loss_fn(X, Y)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

为什么培训需要这么长时间？

Answer 1

为什么 TensorFlow 很慢

TensorFlow 有 2 种执行模式：急切执行和图模式。 TensorFlow 默认行为，从版本 2 开始，默认为急切执行。 Eager Execution 很棒，因为它使您能够编写接近于编写标准 Python 的代码。编写起来更容易，调试起来也更容易。不幸的是，它确实不如图形模式快。

所以这个想法是，一旦函数在 Eager 模式下原型化，让 TensorFlow 在图形模式下执行它。为此，您可以使用 tf.function。 tf.function 将可调用对象编译为 TensorFlow 图。一旦函数被编译成图形，性能增益通常是非常重要的。在 TensorFlow 中开发时推荐的方法如下：

<块引用>

• 在 Eager 模式下调试，然后用 @tf.function 装饰。
• 不要依赖 Python 的副作用，例如对象突变或列表追加。
• tf.function 与 TensorFlow ops 配合使用效果最佳； NumPy 和 Python 调用被转换为常量。

我想补充一点：想想你的程序的关键部分，哪些应该首先转换成图形模式。它通常是您调用模型以获得结果的部分。在这里，您会看到最好的改进。

您可以在以下指南中找到更多信息：

• Better performance with tf.function
• Introduction to graphs and tf.function

将 tf.function 应用于您的代码

因此，您至少可以在代码中更改两件事以使其运行得更快：

1. 第一个是不要在少量数据上使用 model.predict。该函数用于处理庞大的数据集或生成器。 （见this comment on Github）。相反，您应该直接调用模型，为了提高性能，您可以将模型调用包装在 tf.function 中。

<块引用>

Model.predict 是一个顶级 API，设计用于在任何循环之外进行批量预测，具有 Keras API 的全部功能。

2. 第二个是使您的训练步骤成为一个单独的函数，并使用 @tf.function 修饰该函数。

因此，我会在您的训练循环之前声明以下内容：

# to call instead of model.predict
model_func = tf.function(model)

def get_train_func(model, model2, loss_fn, optimizer):
    """Wrapper that creates a train step using the two model passed"""
    @tf.function
    def train_func(state1_batch, state2_batch, done_batch, reward_batch, action_batch):
        Q2 = model2(state2_batch) #B
        with tf.GradientTape() as tape:
            Q1 = model(state1_batch)
            Y = reward_batch + gamma * ((1-done_batch) * tf.math.reduce_max(Q2, axis=1))
            # gather is more efficient than a list comprehension, and needed in a tf.function
            X = tf.gather(Q1, action_batch, batch_dims=1)
            loss = loss_fn(X, Y)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        return loss
    return train_func

# train step is a callable 
train_step = get_train_func(model, model2, loss_fn, optimizer)

您可以在训练循环中使用该函数：

if len(replay) > batch_size:
    minibatch = random.sample(replay, batch_size)
    state1_batch = np.array([s1 for (s1,a,r,s2,d) in minibatch]).reshape((batch_size, 64))
    action_batch = np.array([a for (s1,a,r,s2,d) in minibatch])   #TODO: Posibles diferencies
    reward_batch = np.float32([r for (s1,a,r,s2,d) in minibatch])
    state2_batch = np.array([s2 for (s1,a,r,s2,d) in minibatch]).reshape((batch_size, 64))
    done_batch = np.array([d for (s1,a,r,s2,d) in minibatch]).astype(np.float32)

    loss = train_step(state1_batch, state2_batch, done_batch, reward_batch, action_batch)
    losses.append(loss)

您还可以进行其他更改以使您的代码更加TensorFlowesque，但是通过这些修改，您的代码在我的 CPU 上需要大约 2 分钟的时间。 （赢率为 97%）。