大致问题是这样的
模型先计算一个输出
然后根据这个输出,用lbfgs去优化另一个变量
最后优化模型,大致代码如下
optimizer = optim.Adam(model.parameters(), lr=lr)
for inputs in dataloader:
outputs = model(inputs)
u = outputs.reshape(1, -1, 1)
beta = torch.zeros(1, n, 1, device=device, requires_grad=True)
lbfgs = torch.optim.LBFGS([beta], lr=lr_lbfgs, line_search_fn='strong_wolfe')
def closure():
lbfgs.zero_grad()
# max g => min -g
g = -G(C, beta, u)
g.backward(retain_graph=True)
return g
lbfgs.step(closure)
loss = G(C, beta, u, v, self.blur)
optimizer.zero_grad()
# minimize W by optimizing theta
loss.backward()
optimizer.step()
但是就会发现,模型一轮要跑15分钟,非常慢
原因:应该是模型在lbfgs那里算了整个模型的梯度而不是只算beta的梯度
解决方法:
查阅文档,发现backward还有一个inputs的参数,于是
可以将代码改成
optimizer = optim.Adam(model.parameters(), lr=lr)
for inputs in dataloader:
outputs = model(inputs)
u = outputs.reshape(1, -1, 1)
beta = torch.zeros(1, n, 1, device=device, requires_grad=True)
lbfgs = torch.optim.LBFGS([beta], lr=lr_lbfgs, line_search_fn='strong_wolfe')
def closure():
lbfgs.zero_grad()
# max g => min -g
g = -G(C, beta, u)
g.backward(inputs=beta)
return g
lbfgs.step(closure)
beta = beta.detach()
beta.requires_grad = False
loss = G(C, beta, u, v, self.blur)
optimizer.zero_grad()
# minimize W by optimizing theta
loss.backward()
optimizer.step()