El problema general es este.
El modelo primero calcula una salida
y luego usa lbfgs para optimizar otra variable basada en esta salida.
Finalmente, el modelo se optimiza. El código aproximado es el siguiente
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()
Pero encontrará que el modelo tarda 15 minutos en ejecutarse durante una ronda, lo cual es muy lento.
Motivo: Debería ser que el modelo calcule el gradiente de todo el modelo en lbfgs en lugar de solo el gradiente de beta.
Solución:
consulte la documentación y descubra que hacia atrás también tiene un parámetro de entrada, por lo que
El código se puede cambiar a
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()