python - 如何在 autograd 反向传播中禁用某些模块的梯度更新？

Question

我正在构建一个用于强化学习的多模型神经网络，其中包括一个 Action 网络、一个世界模型网络和一个评论家。这个想法是训练世界模型根据来自 Action 网络的输入和之前的状态来模拟你试图掌握的任何模拟，训练评论家根据世界模型输出最大化贝尔曼方程(随时间的总强化)，然后通过世界模型反向传播批评值，为训练 Action 提供梯度目标。因此 - 从某个状态开始， Action 网络输出一个 Action ，该 Action 被馈送到模型中以生成下一个状态，并且该状态被馈送到批评网络以针对某个目标状态进行评估。

为了使所有这些都起作用，我必须使用 3 个单独的损失函数，每个网络一个，它们都会为一个或多个网络中的梯度添加一些东西，但它们可能会发生冲突。例如 - 为了训练世界模型，我使用来自环境模拟的目标，对于评论家，我使用当前状态奖励 + 折扣 * 下一状态预测值的目标。然而，为了训练 a actor，我只是使用负评值(value)作为损失，并通过所有三个模型一路反向传播以校准最佳 Action 。

我可以在没有任何批处理的情况下通过渐进地清零梯度来完成这项工作，但这样做效率低下，并且不允许我为任何类型的“时间序列批处理”优化器更新步骤累积梯度。每个模型都有自己的可训练参数，但执行图流经所有三个网络。因此在按顺序触发网络后的校准循环中:...

        if self.actor.calibrating:
            self.actor.optimizer.zero_grad()
            #Pick loss For maximizing the value of all actions
            loss = -self.critic.value
            #Backpropagate through all three networks to train actor output
            #How do I stop the critic and model networks from incrementing their gradient values?
            loss.backward(retain_graph=True)
            self.actor.optimizer.step()
        if self.model.calibrating:
            self.model.optimizer.zero_grad()
            #Reduce loss for ambiguous actions
            loss = self.model.get_loss() * self.actor.get_confidence()**2
            #How can I block this from backpropagating through action network?
            loss.backward(retain_graph=True)
            self.model.optimizer.step()
        if self.critic.calibrating:
            self.critic.optimizer.zero_grad()
            #Reduce loss for ambiguous actions
            loss = self.critic.get_loss(self.goal) * self.actor.get_confidence()**2
            #How do I stop this from backpropagating through the model and action networks?
            loss.backward(retain_graph=True)
            self.critic.optimizer.step()

...

最后 - 我的问题分为两部分:

1. 如何在给定层暂时停止 loss.backward() 而不永远分离它？
2. 如何阻止 loss.backward() 更新一些梯度，而我只是在流过一个模型以获得另一个模型的梯度？

Answer 1

由于同事的建议尝试 requires_grad 设置，解决了这个问题。 (我原以为这会破坏执行图，但事实并非如此)

所以 - 回答我自己的两个问题:

1. 如果您以正确的顺序校准链接模型，您可以一次将它们分离，这样 loss.backward() 就不会运行不需要的模型。我在想这会破坏图表，但是......这是 Pytorch，而不是 Tensorflow 1.x，并且无论如何都会在每次前向传递时重新生成图表。傻我昨天错过了这个。
2. 如果您将模型(或图层或单个权重)的 requires_grad 设置为 False，则 loss.backward() 仍将遍历整个连接图，但它会保留那些单独的梯度，同时仍会更早地设置任何梯度在图中。正是我想要的。

此代码可最大程度地减少不必要的图形遍历和梯度更新的执行。我仍然需要重构它以随着时间的推移进行交错更新，以便它可以在逐步优化器之前累积几个循环的梯度，但这绝对可以按预期工作。

#Step through all models in a chain to create gradient paths from critic back through the world model, to the actor.
    def step(self):
        #Get the current state from the simulation
        state = self.world.state
        #Fire the actor to select a softmax action.
        self.actor(state)
        #run the world simulation on that action.
        self.world.step(self.actor.action)
        #Combine the action and starting state as input to the world model.
        if self.actor.calibrating:
            action_state = torch.cat([self.actor.value, state], dim=0)
        else:
            #Push softmax action closer to 1.0
            action_state = torch.cat([self.actor.hard_value, state], dim=0)
        #Run the model and then the critic on the action_state
        self.critic(self.model(action_state))
        if self.actor.calibrating:
            self.actor.optimizer.zero_grad()
            self.model.requires_grad = False
            self.critic.requires_grad = False
            #Pick loss For maximizing the value of the action choice
            loss = -self.critic.value * self.actor.get_confidence()
            loss.backward(retain_graph=True)
            self.actor.optimizer.step()
        if self.model.calibrating:
            #Don't need to backpropagate through actor again
            self.actor.value.detach_()
            self.model.optimizer.zero_grad()
            self.model.requires_grad = True
            #Reduce loss for ambiguous actions
            loss = self.model.get_loss() * self.actor.get_confidence()**2
            loss.backward(retain_graph=True)
            self.model.optimizer.step()
        if self.critic.calibrating:
            #Don't need to backpropagate through the model or actor again
            self.model.value.detach_()
            self.critic.optimizer.zero_grad()
            self.critic.requires_grad = True
            #Reduce loss for ambiguous actions
            loss = self.critic.get_loss(self.goal) * self.actor.get_confidence()**2
            loss.backward(retain_graph=True)
            self.critic.optimizer.step()