tensorflow的autodiff比pytorch的对应项慢

Question

我正在使用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()

Answer 1

def __init__(self,...):
  ...
  self.model.call = tf.function(self.model.call)
  ...

您需要使用tf.function来包装模型的调用函数。