Optimizing arbitrary parameter using Adam optimizer from PyTorch

optimize.py

import torch
import torch.nn as nn

LEARNING_RATE = 3e-4

# we consider a toy example: minimize the squared value of single vector.
x = torch.randn(5)
m = nn.MSELoss() # for calculating loss
optimizer = torch.optim.Adam([x.requires_grad_()], lr=LEARNING_RATE)

# through this iteration, we can observe that the values (slowly) converge to 0.
# use larger LEARNING_RATE if you want a faster convergence.
for step in range(10):
  print(x)
  loss = m(x, torch.zeros(5))
  loss.backward()
  optimizer.step()
  optimizer.zero_grad()

Example output:

tensor([-1.4675, -1.6058, -0.8570, -1.4888,  0.6637], requires_grad=True)
tensor([-1.4672, -1.6055, -0.8567, -1.4885,  0.6634], requires_grad=True)
tensor([-1.4669, -1.6052, -0.8564, -1.4882,  0.6631], requires_grad=True)
tensor([-1.4666, -1.6049, -0.8561, -1.4879,  0.6628], requires_grad=True)
tensor([-1.4663, -1.6046, -0.8558, -1.4876,  0.6625], requires_grad=True)
tensor([-1.4660, -1.6043, -0.8555, -1.4873,  0.6622], requires_grad=True)
tensor([-1.4657, -1.6040, -0.8552, -1.4870,  0.6619], requires_grad=True)
tensor([-1.4654, -1.6037, -0.8549, -1.4867,  0.6616], requires_grad=True)
tensor([-1.4651, -1.6034, -0.8546, -1.4864,  0.6613], requires_grad=True)
tensor([-1.4648, -1.6031, -0.8543, -1.4861,  0.6610], requires_grad=True)