Basics

Domain experts frequently use machine learning in their fields without needing in-depth knowledge of a computer’s inner workings. Rewriting their scripts for performance gains is challenging typically requires expertise different from that of the original author. SOL is designed for these domain experts and aims to optimize machine learning models without the need to understand the underlying hardware. To this end SOL is designed with two main principles in mind:

  • Ease of use: You do not need to select any SOL specific parameters. All information regarding execution parameters is read directly from the model.

  • Drop-in replacement: You do not have to rewrite any of your code. Optimize your model with SOL and use the generated model in place of your old one.

To optimize a model with SOL you just need to add

import sol

to your imports and call

optimized_model = sol.optimize(model)

on your framework model.

And that’s it, you now have an optimized version of your previous model!

According to the design goal of drop-in replacement, sol.optimize creates a model of the same type as its input. This means that type(optimized_model) will be equal to type(model). This means for example an nn.Module for torch or tf.keras.Model for Tensorflow. This includes any custom functions you have defined for your model, e.g., model.do_what_I_want(...) are also preserved (but not optimized by SOL). As a result, you do not need to change anything else in your codebase and just replace the old model with the optimized one.

SOL uses JIT-compilation (just in time) for its final result. This means that compilation is not triggered in sol.optimize(), but when the model is actually called the first time.

If you already use torch.compile in your project, adding SOL is even easier. You just have to add backend="sol" to your torch.compile() call.

import torch
import torchvision.models as models
import sol.pytorch # not needed for torch >= 2.6.0

model = models.resnet18(pretrained=False)
optimized_model = torch.compile(model, backend="sol")

If you are using a pytorch version older than 2.6 you also need to add import sol.pytorch. For any version past 2.6, SOL is automatically detected and added to valid backends during installation.

Running Inference

The optimized model behaves exactly the same as the framework model beforehand. To use a SOL-model you just replace your old torch model with the optimized one like this:

import torch
import torchvision.models as models
import sol

model = models.resnet18()
model.eval()

optimized_model = sol.optimize(model)

# Generate a random input tensor 
random_input = torch.randn(1, 3, 224, 224)
# Run Inference
with torch.no_grad():
    # out = model(random_input)
    out = optimized_model(random_input)

Training the Model

Now, let’s train the model. For this will use a simple example training a small MLP on FashionMNIST. Training requires some more setup, so we define helpers like a dataloader and a definition of the network:

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# Define dataloader
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)
dataloader = DataLoader(training_data, batch_size=64)

# Define model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork().train()

# Define the Loss Function
loss_fn = nn.CrossEntropyLoss()

# Optimize model
import sol
model = sol.optimize(model)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) 

# Training loop
num_epochs = 5
for epoch in range(num_epochs):
    for batch, (X, y) in enumerate(train_dataloader):
        # Compute prediction error
        pred = model(X) 
        loss = loss_fn(pred, y)

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

# Save your model to disk
torch.save(model.state_dict(), "sol_model")

As before with inference, the only change needed to compile the model with SOL is the call to sol.optimize on the model before it is used in the training process. This example showcases again how the optimized model can be used as a drop-in replacement for the original one. Not only is it used to compute the prediction error, model.parameters() is also used to initialize the optimizer and torch.save saves the models parameters without any changes to equivalent pure PyTorch code.

This feature also allows you write your scripts with SOL as an optional component. You can simply wrap the appropriate code in a try block or make it dependent on some parameter of your script. This way you can reuse your script on any machine even if it does not have SOL installed. It also makes sure that any problems that may occur during SOL’s optimization will never be a point of critical failure to your script.

if use_sol:
    try:
        import sol
        model = sol.optimize(model)
    except:
        pass