How to use Pytorch as a general optimizer | by Conor Mc. | Towards Data Science

Pytorch is really fun to work with and if you are looking for a framework to get started with neural networks I highly recommend it — see my short tutorial on how to get up and running with a basic neural net in Pytorch here.
What many people don’t realise however is that Pytorch can be used for general gradient optimization. In other words, you can use Pytorch to find the minium or maximum of arbitrarily complex optimization objectives. But, why would you want to do this? I can think of at least three good reasons (there are many more).
You are already familiar with Pytorch and don’t want to have to learn another optimization frameworkYou want to optimize over the outcomes of a Pytorch model — i.e. you want to use optimize over the predictions of a Pytorch Neural net (e.g. a first stage neural net might predict the propensity of a customer to engage in a particular high-value action and the optimizer is used to determine which action is best given some contraints such as marketing budget).You want to use the advanced optimizers defined in Pytorch such as Adam.
You are already familiar with Pytorch and don’t want to have to learn another optimization framework
You want to optimize over the outcomes of a Pytorch model — i.e. you want to use optimize over the predictions of a Pytorch Neural net (e.g. a first stage neural net might predict the propensity of a customer to engage in a particular high-value action and the optimizer is used to determine which action is best given some contraints such as marketing budget).
You want to use the advanced optimizers defined in Pytorch such as Adam.
Well ... you don’t actually have to implement anything, if you are familiar with Pytorch already you simply write a Pytorch custom module in the same way you would for a neural network and Pytorch will take care of everything else. Let’s see a worked example.
To demonstate we’ll define a simple function — an expoential decay function. Let’s define the data — the scaler floats a, k, b are the unknown parameters of the function and the objective of the optimization will be to estimate these parameters.
import numpy as npimport torchimport matplotlib.pyplot as pltfrom torch import nnfrom torch.functional import Ffrom copy import copyimport seaborn as snssns.set_style("whitegrid")n = 1000noise = torch.Tensor(np.random.normal(0, 0.02, size=n))x = torch.arange(n)a, k, b = 0.7, .01, 0.2y = a * np.exp(-k * x) + b + noiseplt.figure(figsize=(14, 7))plt.scatter(x, y, alpha=0.4)
Next, let’s define the model and the training loop:
class Model(nn.Module):    """Custom Pytorch model for gradient optimization.    """    def __init__(self):                super().__init__()        # initialize weights with random numbers        weights = torch.distributions.Uniform(0, 0.1).sample((3,))        # make weights torch parameters        self.weights = nn.Parameter(weights)                    def forward(self, X):        """Implement function to be optimised. In this case, an exponential decay        function (a + exp(-k * X) + b),        """        a, k, b = self.weights        return a * torch.exp(-k * X) + b    def training_loop(model, optimizer, n=1000):    "Training loop for torch model."    losses = []    for i in range(n):        preds = model(x)        loss = F.mse_loss(preds, y).sqrt()        loss.backward()        optimizer.step()        optimizer.zero_grad()        losses.append(loss)      return losses
If you are familiar with Pytorch there is nothing too fancy going on here. The key thing that we are doing here is defining our own weights and manually registering these as Pytorch parameters — that is what these lines do:
weights = torch.distributions.Uniform(0, 0.1).sample((3,))# make weights torch parametersself.weights = nn.Parameter(weights)
The lines below detemine the function to be optimised. You can replace these with the definition of the function you want to minimise.
a, k, b = self.weightsreturn a * torch.exp(-k * X) + b
By calling nn.Parameter the weight we define will behave and function in the same way as standard Pytorch parameters — i.e they can calculate gradients and be updated in response to a loss function. The training loop is simply iterating over n epochs, each time estimating the mean squared error and updating the gradients.
Time to run the model, we’ll use Adam for the optimization.
# instantiate modelm = Model()# Instantiate optimizeropt = torch.optim.Adam(m.parameters(), lr=0.001)losses = training_loop(m, opt)plt.figure(figsize=(14, 7))plt.plot(losses)print(m.weights)
The plot above shows the loss function over 1000 epochs — you can see that after ~600 it is showing no signs of further improvement. The estimated weights for a, k, b are 0.697, 0.0099, 0.1996, so extremely close to the parameters that define the function and we can use the trained model to estimate the function:
preds = m(x)plt.figure(figsize=(14, 7))plt.scatter(x, preds.detach().numpy())plt.scatter(x, y, alpha=.3)
You can see that with the optimised parameters we can now fit the exponential decay of the data. Whilst you could do the same thing with curve_fit from scipy I think the greater flexibility to fit more complex functions with Pytorch is worth investing the time to learn.
Summary
Whilst most people will use Pytorch for building neural networks, the flexibility of the framework makes it extremely versatile. In this post, we have fit a simple curve defined by the exponential decay function, but there is no reason why the same building blocks cannot be extended to arbitrarily complex optimization functions.
Hope this was helpful, let me know if you have any thoughts, comments or questions below.
Thanks for reading!