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

Question

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

我已经对代码进行了实验，用不到 3 分钟的时间来训练网络并获得 89% 的胜利（训练后赢得 100 场比赛中的 89 场）。

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

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

训练损失数字和获胜率似乎是正确的（我们必须考虑到游戏是随机的并且可能存在不可能的状态）。问题在于整个过程的性能。

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

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

        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 的代码。它更容易编写，也更容易调试。不幸的是，它真的没有图形模式快。

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

• 在 Eager 模式下调试，然后用 @tf.function 装饰。
• 不要依赖 Python 的副作用，例如对象突变或列表追加。
• tf.function 最适合 TensorFlow 操作； 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 的全部功能。

1. 第二个是让你的训练步骤成为一个单独的函数，并用@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% 的胜率）。