- 旧版
- 定义一个类,然后用函数的方式去调用:
from torch.autograd import Function
class GradReverse(Function):
def __init__(self, lambd):
self.lambd = lambd
def forward(self, x):
return x.view_as(x)
def backward(self, grad_output):
return (grad_output * -self.lambd)
def grad_reverse(x, lambd=1.0):
return GradReverse(lambd)(x)
- 新版
class GradReverse(Function):
# 重写父类方法的时候,最好添加默认参数,不然会有warning(为了好看。。)
@ staticmethod
def forward(ctx, x, lambd, **kwargs: None):
# 其实就是传入dict{'lambd' = lambd}
ctx.lambd = lambd
return x.view_as(x)
@staticmethod
def backward(ctx, *grad_output):
# 传入的是tuple,我们只需要第一个
return grad_output[0] * -ctx.lambd, None
# 这样写是没有warning,看起来很舒服,但是显然是多此一举咯,所以也可以改写成
@staticmethod
def backward(ctx, grad_output):
# 直接传入一格数
return grad_output * -ctx.lambd, None
新旧区别:
(1)无需init,直接转成静态@staticmethod;
(2)在调用时,直接使用类名.apply使用;
(3)需要传参时: ctx.lambda = lambda 第一个参数必须是ctx
注意:
backward报错:Runtime Error: returned an incorrect number of gradients
分析:程序默认返回所有forward每一个输入参数的loss,现在我们并不需要给所有参数返回loss,只需要给input tensor返回就够了,其他的直接返回None