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

Question

我正在通过Deep Reinforcement Learning in Action一书学习 DRL 。在第 3 章中，他们展示了简单游戏 Gridworld（此处的说明，在规则部分）以及PyTorch 中的相应代码。

我已经对代码进行了试验，用 89% 的胜利（在训练后赢得 100 场比赛中的 89 场）训练网络只需不到 3 分钟。

作为练习，我已将代码迁移到tensorflow。所有代码都在这里。

问题是，使用我的 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有两种执行模式：急切执行和图形模式。TensorFlow从版本 2 开始，默认行为是默认立即执行。热切执行非常棒，因为它使您能够编写类似于标准 Python 的代码。它更容易编写，也更容易调试。不幸的是，它确实不如图形模式那么快。

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

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

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

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

• 使用 tf.function 获得更好的性能
• 图和 tf.function 简介

应用于tf.function您的代码

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

1. 第一个是不要使用model.predict少量数据。该函数适用于大型数据集或生成器。（请参阅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%）。