问题描述
我正在重新实现 Invertible Residual Networks 架构。
class iresnetBlock(nn.Module):
def __init__(self,input_size,hidden_size):
self.bottleneck = nn.Sequential(
LinearContraction(input_size,hidden_size),LinearContraction(hidden_size,input_size),nn.ReLU(),)
def forward(self,x):
return x + self.bottleneck(x)
def inverse(self,y):
x = y.clone()
while not converged:
# fixed point iteration
x = y - self.bottleneck(x)
return x
我想向 inverse
函数添加自定义向后传递。由于是不动点迭代,因此可以利用隐函数定理来避免循环展开,而是通过求解线性系统来计算梯度。例如,这是在 Deep Equilibrium Models 架构中完成的。
def inverse(self,y):
with torch.no_grad():
x = y.clone()
while not converged:
# fixed point iteration
x = y - self.bottleneck(x)
return x
def custom_backward_inverse(self,grad_output):
pass
我如何为这个函数注册我的自定义向后传递?我希望这样,当我稍后定义一些损失(例如 r = loss(y,model.inverse(other_model(model(x))))
)时,r.backwards()
会正确使用我的自定义梯度进行反向调用。
理想情况下,解决方案应与 torchscript
兼容。
解决方法
暂无找到可以解决该程序问题的有效方法,小编努力寻找整理中!
如果你已经找到好的解决方法,欢迎将解决方案带上本链接一起发送给小编。
小编邮箱:dio#foxmail.com (将#修改为@)