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

Posted 2023-02-23

【中文标题】为啥这个 tensorflow 训练需要这么长时间？

【英文标题】：Why is this tensorflow training taking so long?为什么这个 tensorflow 训练需要这么长时间？

【发布时间】：2021-07-26 17:20:44

【问题描述】：

我正在通过深度强化学习实战一书学习 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))

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

【问题讨论】：

您可以尝试禁用急切执行并查看它的效果吗？ tf.compat.v1.disable_eager_execution()

我还没有完成循环，因为它花费了太多时间：5000 个 epoch 中的 228 个需要 14 分钟。所以，它似乎比没有 disable_eager 慢一点。

唯一明显的区别是，没有disable_eager_execution，CPU 大约是 30-35%。使用disable_eager_execution 几乎达到 100%，界面感觉很慢。

嗨@Gulzar，两个代码（PyTorch 和 Tensorflow）都在问题中链接

尝试明确传递training=False or True。 Q2 = model2(state2_batch, training=False) 和 Q1 = model(state1_batch, training=True)

【参考方案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 应用于您的代码

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

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

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

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

【讨论】：

好的，我回家试试，但是..为什么？为什么这些变化会从小时变为分钟？在第一次更改的情况下，我可以理解一个函数针对少量数据进行了优化，而另一个针对大量数据进行了优化，但是第二次更改如何提高执行时间？

这很好。我添加了一个解释以及一些指向指南的链接，这些链接应该可以帮助您更好地理解为什么。

答案很好，我想补充一点，profile tensorflow using tensorboard有选项

@Lescurel 毫无疑问，这是一个非常有见地的答案。但我认为这更像是优化 TF 模型的方式。与图形模式相比，Eager 模式很慢，但在 OP 的方法中，TF 和 PyTorch 都处于 Eager 模式，但存在显着的执行差距。我认为这是重点；有什么指针吗？

@M.Innat 我不知道这是否是一个错误。我是TF的初学者。如果您对它的错误有足够的信心，请不要犹豫，告知它。