pytorch GradScale 梯度缩放算子

最新推荐文章于 2025-11-25 18:10:22 发布

原创最新推荐文章于 2025-11-25 18:10:22 发布 · 3.1k 阅读

3 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#pytorch #梯度缩放

深度学习专栏收录该内容

47 篇文章

订阅专栏

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

Qwen3-32B-Chat 私有部署镜像 | RTX4090D 24G 显存 CUDA12.4 优化版

Qwen

文本生成

Qwen3

本镜像基于 RTX 4090D 24GB 显存 + CUDA 12.4 + 驱动 550.90.07 深度优化，内置完整运行环境与 Qwen3-32B 模型依赖，开箱即用。

很早前就想弄这东西，正好有空。
这个东西可以很方便的缩放指定层或指定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())