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

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

1. 使用 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())

2. 写一个反传算子来实现

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())