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

Question

我正在用这本书学习 DRL 深度强化学习实战 .在第 3 章中，他们介绍了简单的游戏 Gridworld(instructions here，在规则部分)，相应的代码在 中。 PyTorch .
我已经对代码进行了试验，用不到 3 分钟的时间来训练网络并获得 89% 的胜利(训练后赢得 100 场比赛中的 89 场)。

作为练习，我已将代码迁移到 tensorflow .所有代码都是here .
问题是，使用我的 tensorflow 端口，训练网络需要将近 2 个小时，胜率达到 84%。两个版本都使用唯一的 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 模式下原型(prototype)化，让 TensorFlow 以图形模式执行它。为此，您可以使用 tf.function . tf.function将可调用对象编译为 TensorFlow 图。一旦函数被编译成图表，性能增益通常是非常重要的。在 TensorFlow 中开发时推荐的方法如下:

• Debug in eager mode, then decorate with @tf.function.
• Don't rely on Python side effects like object mutation or list appends.
• tf.function works best with TensorFlow ops; NumPy and Python calls are converted to constants.

我要补充一点:考虑一下程序的关键部分，哪些应该首先转换为图形模式。它通常是您调用模型以获得结果的部分。这是您将看到最好的改进的地方。
您可以在以下指南中找到更多信息:

Better performance with tf.function

Introduction to graphs and tf.function

申请tf.function到你的代码
因此，您至少可以在代码中更改两件事以使其运行得更快:

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

Model.predict is a top-level API designed for batch-predicting outside of any loops, with the fully-features of the Keras APIs.

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