Python tensorflow autodiff比pytorch慢'；中国同行

我正在使用tensorflow 2.0并尝试评估梯度，以便将其反向传播到一个简单的前馈神经网络。以下是我的模型的外观：

def __init__(self, input_size, output_size):
        inputs = tf.keras.Input(shape=(input_size,))
        hidden_layer1 = tf.keras.layers.Dense(30, activation='relu')(inputs)
        outputs = tf.keras.layers.Dense(output_size)(hidden_layer1)
        self.model = tf.keras.Model(inputs=inputs, outputs=outputs)

        self.optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
        self.loss_function = tf.keras.losses.Huber()

这个网络的前向传递很好，但当我使用梯度带训练模型时，它至少比PyTorch慢10倍。 培训职能：

def learn_modified_x(self, inputs, targets, actions):
        with tf.GradientTape() as tape:
            predictions = self.model(inputs)
            predictions_for_action = gather_single_along_axis(predictions, actions)
            loss = self.loss_function(targets, predictions_for_action)

        grads = tape.gradient(loss, self.model.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.model.trainable_weights))

我试着用注释的方式来找出问题的真正原因。我发现tape.gradient是造成这种情况的一个重要因素

有什么想法吗

PyTorch实现

    def __init__(self, input_size, nb_action):
        super(Network, self).__init__()
        self.input_size = input_size
        self.nb_action = nb_action
        self.fc1 = nn.Linear(input_size, 30)
        self.fc2 = nn.Linear(30, nb_action)
    
    def forward(self, state):
        x = F.relu(self.fc1(state))
        q_values = self.fc2(x)
        return q_values

    def learn(self, batch_state, batch_next_state, batch_reward, batch_action):
        outputs = self.model(batch_state).gather(1, batch_action.unsqueeze(1)).squeeze(1)
        next_outputs = self.model(batch_next_state).detach().max(1)[0]
        target = self.gamma*next_outputs + batch_reward
        td_loss = F.smooth_l1_loss(outputs, target)
        self.optimizer.zero_grad()
        td_loss.backward(retain_variables = True)
        self.optimizer.step()

def\uuuuu初始化（self，…）：
...
self.model.call=tf.function（self.model.call）
...

---

您需要使用

tf.function

来包装您的模型的调用函数。

我想性能有了一些提升，但与PyTorch实现相比仍然非常缓慢