40 lines
792 B
Python
40 lines
792 B
Python
import torch
|
|
|
|
def conjugate_gradient(A, b, max_iters, res_tol=1e-10):
|
|
x = torch.zeros_like(b)
|
|
r = b - A(x)
|
|
p = r
|
|
|
|
rTr = r.T @ r
|
|
|
|
for _ in range(max_iters):
|
|
Ap = A(p)
|
|
alpha = rTr / (p.T @ Ap)
|
|
x = x + alpha * p
|
|
|
|
r = r - alpha * Ap
|
|
if torch.norm(r) < res_tol:
|
|
break
|
|
|
|
rTrnew = r.T @ r
|
|
beta = rTrnew / rTr
|
|
p = r + beta * p
|
|
rTr = rTrnew
|
|
|
|
return x
|
|
|
|
def line_search(f, x0, dx, g0, alpha, condition, max_steps=10, c1=0.1):
|
|
assert 0 < alpha < 1
|
|
|
|
f0 = f(x0)
|
|
for _ in range(max_steps):
|
|
x = x0 + dx
|
|
|
|
if (f(x) > f0 + c1 * g0.T @ dx) and condition(x):
|
|
return x
|
|
|
|
dx *= alpha
|
|
|
|
print('Line search failed, returning x0')
|
|
return x0
|