pytorch GradScale 梯度缩放算子

最新推荐文章于 2025-10-22 14:20:34 发布

原创最新推荐文章于 2025-10-22 14:20:34 发布 · 3.1k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #梯度缩放

深度学习专栏收录该内容

44 篇文章

订阅专栏

本文介绍了如何在PyTorch中利用Tensor.register_hook和自定义GradScaleOp类实现梯度缩放，通过实例展示了如何在StyleGAN的MapNet中类似权重缩放的方式调整学习率。两种方法分别展示了动态调整和内嵌操作的灵活性。

部署运行你感兴趣的模型镜像

很早前就想弄这东西，正好有空。
这个东西可以很方便的缩放指定层或指定Tensor的梯度。与调整对应变量的学习率一个效果。
Stylegan 里面 MapNet 使用了另外一种基于权重缩放的学习率缩放方式，算了，效果差不多。

有两个方便的方法可以实现。

使用 Tensor.register_hook 来调整

import torch

x = torch.ones(1, dtype=torch.float32, requires_grad=True)

y = x + 0
y.backward()
# 可以看到 x.grad 值为1
print(x.item(), y.item(), x.grad.item())

del x.grad

y = x + 0
y.register_hook(lambda g: g*0.1)
y.backward()
# 可以看到 x.grad 值为0.1，成功缩放梯度
print(x.item(), y.item(), x.grad.item())

写一个反传算子来实现

import torch
import torch.autograd


class GradScaleOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale):
        scale = torch.as_tensor(scale, dtype=x.dtype, device=x.device)
        ctx.save_for_backward(scale)
        return x

    @staticmethod
    def backward(ctx, grad_output):
        s, = ctx.saved_tensors
        return grad_output * s, None

gradscale = GradScaleOp.apply

x = torch.ones(1, dtype=torch.float32, requires_grad=True)

y = x + 0
y.backward()
# 可以看到 x.grad 值为1
print(x.item(), y.item(), x.grad.item())

del x.grad

y = x + 0
y = gradscale(y, 0.1)
y.backward()
# 可以看到 x.grad 值为0.1，成功缩放梯度
print(x.item(), y.item(), x.grad.item())