Pytorch对梯度进行rescale

继承torch.autograd.Function类，然后定义一个前向和反向函数，ctx是自动求导过程中用于存储属性的对象，比如ctx.save_for_backward。在foward中可以定义计算方程，但是本例子中只是对某一层的梯度进行rescale，所以对输入不做处理。

class GradientRescaleFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, weight):
        ctx.save_for_backward(input)
        ctx.gd_scale_weight = weight
        output = input
        return output
        
    @staticmethod
    '''backward中的输入是grad_output，然后值得注意的是，前传进来几个参数，反传的梯度就要有几个，因此，return了input 和weight梯度，由于weight只是一个常数，所以给其梯度置为None'''
    def backward(ctx, grad_output):
        input = ctx.saved_tensors
        grad_input = grad_weight = None

        if ctx.needs_input_grad[0]:
            grad_input = ctx.gd_scale_weight * grad_output

        return grad_input, grad_weight
    
gradient_rescale = GradientRescaleFunction.apply