How to use PyTorch as a general function optimizer?

I have a parametric function that should model the behavior of some data Y, at input position X. I want to use pytorch optimizers and a GPU, but the tutorials out there assume that I want to use neural layers. Would you help me defining a minimum working example?

I tried following the official guide of PyTorch, but I could not get a minimal working example that do not use schedulers or neural networks. I had a look for similar answers on stackoverflow, but they assume a knowledge that I do not have - I do have knowledge on which are the steps to be done for an optimization with a numpy/scipy combination though

Solution

To begin, create torch tensors (similar to numpy arrays) for your X and Y data, and ensure they're sent to the GPU. In this example, we'll use the first CUDA GPU to store a single input value "2." and a desired output "10.":

import torch
device = "cuda:0"
X = torch.tensor(2., device=device)
Y = torch.tensor(10., device=device)

Next, define your model. The model may consist of various components, conceptualized as neural layers or an ensemble of parameters, usually stated in the initialization. In this case, we have a single parameter (initially set to 0) representing the scaling factor in a parametric function. The function scales the input and is defined with the forward method.

class Model(torch.nn.Module):
  def __init__(self, device):
    super(Model, self).__init__()
    init_scaling_factor = torch.zeros(1, device=device)
    self.scaling_factor = torch.nn.Parameter(init_scaling_factor, requires_grad=True)
 
  def forward(self, x):
    output = self.scaling_factor * x
    return output

Now, initialize the model, an SGD optimizer, and a cost function.

model = Model(device=device)
optim = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = torch.nn.MSELoss()

In the optimization loop, calculate the cost function at each step. Use the optimizer to minimize it through the loss.backward() and optim.step() combination. First, this backpropagates the error through the operations on the input X (so that is known how the optimizer can have a better estimate for the model parameters), and then updates the model parameters.

n_iters = 8000
for i in range(0, n_iters):
  predictions = model.forward(X)
  loss = loss_fn(predictions, Y)
  loss.backward()
  optim.step()
  optim.zero_grad()
  print(f"loss ({i}): {loss.item()}")
print(f"The scaling parameter is {model.scaling_factor.data.item()}")

This loop iteratively refines the model parameters, and at the end, you'll have the optimized scaling factor for your parametric function. Adjust the parameters and structure as needed for your specific use case.