继承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