How to approximate any function using PyTorch

When analyzing data and building machine learning models, there is often a need to approximate complex functions. PyTorch provides convenient tools for creating and training neural networks that can be effectively used for this purpose. In this post, we'll look at a simple example of function approximation using PyTorch.

Step 1: Data Preparation

Let's start with a simple control theory example and solve it first analytically and then using PyTorch. Consider the function

This is a linear inhomogeneous differential equation, where

Exactly time, response we will later use it as a training dataset. We will also depict an analytical solution to the problem.

class Analytic:
    def __init__(self, alpha, beta, c_1, c_2):
        self.alpha = alpha
        self.beta = beta

        self.c_1 = c_1
        self.c_2 = c_2 
    
    def calculate(self, t: float) -> float:
        return self.c_1 * np.exp(self.alpha * t) * np.sin(self.beta * t) \
             + self.c_2 * np.exp(self.alpha * t) * np.cos(self.beta * t)
    
    def __call__(self, t: list[float]) -> list[float]:
        return np.vectorize(self.calculate)

analytic = Analytic(
    alpha = -0.5,
    beta = np.sqrt(3)/2,
    c_1 = 4/np.sqrt(3),
    c_2 = 2
)

plt.plot(time, analytic(time), linestyle="--", 
         label="Analityc solution")

Step 2: Model Definition

Now we need to decide on the architecture of the model. For us this will be a regular function of the form

This code defines the class Modelwhich inherits from the class nn.Module in PyTorch. In method init we define trainable parameters alpha, C1 And C2as well as a fixed frequency beta. In method forward a forward pass of the model is defined, where we calculate the output value output depending on the input variable t using the specified parameters.

To train the model, we need to write a dataset class, a descendant of the class torch.utils.data.Dataset.

class TimeResponseDataset(Dataset):
    """Torch Dataset for x=time and y=response"""
    def __init__(self, time, response):
        self.time = torch.Tensor(time)
        self.response = torch.Tensor(response)
        
    def __len__(self):
        return len(self.time)
    
    def __getitem__(self, idx):
        return {'time': self.time[idx], 'response': self.response[idx]}

We need it so we can throw it in later. torch.utils.data.DataLoader. He will then give us batches of data during training.

dataset = TimeResponseDataset(time, response)
dataloader = DataLoader(dataset, batch_size=512, shuffle=True)

Step 3: Training

Let's see the number of trainable parameters in the model.

model = Model()
print(f"trainable params: {len([i for i in model.parameters()])}")

The result, as we intended, is 3.

Let's initialize the loss function and optimizer. The loss function will tell us how far we are from the optimum, and the optimizer will determine how quickly and in what direction to change the weights to get there.

lr = 0.01
betas = (0.9, 0.999)
history = []

criterion = nn.MSELoss() 
optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=betas, eps=1e-8)

Now let's train the model for 150 epochs, while maintaining the history of the loss function.

epochs = 150
for epoch in tqdm(range(epochs)):
    for i, data in enumerate(dataloader):
        optimizer.zero_grad()
        pred = model(data['time'])
        loss = criterion(pred, data['response'])
        loss.backward()
        optimizer.step()
        
        history.append(loss.detach().numpy())

A graph of the loss function versus the number of training batches is shown. As can be seen from the graph, the model converged at another 500 step. Let's look at the trained parameters and the analytical solution.

Step 4: Comparison

analytic.alpha, analytic.c_1, analytic.c_2

estimation = [float(param.data) for param in model.parameters()]
estimation

Less error . Let's look at how well our model fits the real data.

approx = Analytic(
    alpha=estimation[0],
    beta=model.beta, 
    c_1=estimation[1],
    c_2=estimation[2]
)

# Analytic, we defined earlier
plt.plot(time, analytic(time), color[0], 
        label="$y=4/\sqrt{3} \cdot e^{-t/2} sin(\sqrt{3}t/2) - 2 e^{-t/2} cos(\sqrt{3}t/2)$.")

# Approximation
plt.plot(time, approx(time), color[2], linestyle="--", label="Approximation")

To understand how quickly the model converges on this problem, you can look at the graphs below. This is a comparison of the analytical solution and our model on epochs 0, 10, 20 .

Additionally

As you may have noticed, not all parameters in the model are trainable. Coefficient that should have been included in periodic functions was immediately recorded. Why was this done? I'm showing you.

I'll explain. Cosine and sine, which are in our approximating model