在训练时发现梯度范数一直是0

Question

def forward(self, x):
    with torch.no_grad():
        x = self.vae.encode(x).latent_dist.sample().mul_(0.15)
    for layer in self.layers:
        x = layer(x)
    with torch.no_grad():
        x = self.vae.decode(x / 0.18215).sample
    x.requires_grad_(True)
    return x

以上是一个forward函数,使用VAE来进行encode和decode，但是我发现训练时梯度范数grad_norm从一开始就是0,但是损失又在慢慢下降，这是为什么？如果不加上VAE，梯度范数就是正常的，如下:

def forward(self, x):
    for layer in self.layers:
        x = layer(x)
    return x

Answer 1

🌈🌈🌈参考通义千问和郭老师的小迷弟雅思莫了-编写提供🌈🌈🌈
您可以参考如下，如果回答的不正确，及时评论区回复我，我会根据你错误描述追加回复，直到您满意为止。

观察你的代码，问题可能是因为你在前向传播的过程中使用了torch.no_grad()。torch.no_grad()的作用是暂时禁止对tensor数据进行梯度追踪。在这个上下文管理器中完成的所有运算都不会被追踪。 在你的代码中，你对VAE的编码和解码过程使用了torch.no_grad()，这意味着这两个步骤对于梯度计算是不可见的，即使你在后面对x使用了requires_grad_(True)，那也只能对x进行梯度追踪，而不能恢复你在torch.no_grad()中丢失的梯度信息。 如果在训练过程中，你对这两个步骤的梯度感兴趣，那么需要移除torch.no_grad()。不过，请注意，如果你同时更新了VAE和其他层，可能需要为VAE和其他层设置不同的学习率，因为通常在微调预训练模型时，直接使用预训练模型的学习率会导致模型快速收敛然后反弹。 可以参考这样改：

def forward(self, x):
    # 移除torch.no_grad()
    x = self.vae.encode(x).latent_dist.sample().mul_(0.15)
    for layer in self.layers:
        x = layer(x)
    x = self.vae.decode(x / 0.18215).sample
    x.requires_grad_(True)
    return x

以上修改只是移除了torch.no_grad()，如果你仍然希望在特定部分禁用梯度计算，但又需要梯度信息，你可能需要手动保存特定tensor的梯度，然后在禁用梯度计算之后手动恢复。具体做法会比较复杂，需要具体任务进行设计。