Implementing a simple optimization algorithm in PyTorch

Question

I'm currently learning PyTorch in order to utilize its open source autograd feature, and as an exercise for myself, I want to implement a simple optimization algorithm that I've already implemented in MATLAB. As a simple example, say I'm trying to solve the problem min_x 1/2 x'Ax - b'x, i.e. find the vector x which minimizes the quantity x'Ax - b'x, given both A and b. A simple gradient descent algorithm with exact line search in MATLAB may look something like this:

% initialize x = zeros(n, 1) where n is the length of b
while residual > tolerance
    grad = A*x - b; % compute the gradient of the objective
    alpha = norm(grad)^2/(grad'*A*grad); % compute step-size alpha by exact line search
    x = x - alpha*grad; % do a gradient step
    residual = norm(grad); % compute residual
    objective = x'*A*x - b'*x; % compute objective value at current iteration

How would I implement this optimization algorithm in PyTorch? Specifically, I would want to do an identical optimization loop, where I replace my own computation of the gradient with Torch's autograd feature. In other words, I want to perform the exact same algorithm as above in PyTorch, except instead of computing the gradient myself, I simply use PyTorch's autograd feature to compute the gradient. In this way, I don't want to call any given optimizer (like SGD or Adam)--just write the algorithm myself with the only difference being that the gradient is computed by PyTorch. I plan to compare the results of the above numpy/MATLAB implementation with the explicit gradient computation vs the PyTorch implementation with what I assume is a numerical approximation of the gradient.

Thank you very much!

ihdv · Accepted Answer · 2020-12-02 12:56:18Z

3

Basically you need to use the autograd in pytorch. Not a complete program but it'll look something like this along the lines:

In each iteration do the following:

Specify x.requires_grad=True because you need its gradient. Then compute your objective function:

x.requires_grad = True
obj_function = torch.matmul(x.t(),torch.matmul(A,x)) * .5 - torch.matmul(b,x)

Clear old gradient results of x:

x.grad = None

Call backward() on obj_function. This will automatically compute the gradients of the tensors involved in the operation:

obj_function.backward()

Set x.requires_grad to false because you need to modify x, but pytorch does not allow inplace modification of a tensor that requires grad:

x.requires_grad = False
x = x - x.grad * step_size

answered Dec 2, 2020 at 12:56

ihdv

2,3373 gold badges20 silver badges41 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

JVGD Over a year ago

I am surprised you did not have any upvotes, but then I thought I came for looking something way too specific so that may be the reason, anyways, here is my upvote for <3

Collectives™ on Stack Overflow

Implementing a simple optimization algorithm in PyTorch

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related