SymbolicModel Demoο
In this demo, we show you how to:
Wrap an MLP block in a PyTorch model with our
SymbolicModelclassPerform symbolic regresson on the MLP with the
distillmethodSwitch the MLP to an equation in the forward pass of the model with the
switch_to_equationmethodSwitch back to the MLP with the
switch_to_blockmethod
Wrapping a PyTorch modelο
Create a simple PyTorch model.
import torch
import numpy as np
import torch.nn as nn
class MLP(nn.Module):
"""
Simple MLP.
"""
def __init__(self, input_dim, output_dim, hidden_dim):
super(MLP, self).__init__()
self.mlp = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(hidden_dim, output_dim)
)
def forward(self, x):
return self.mlp(x)
class SimpleModel(nn.Module):
"""
Simple model class.
"""
def __init__(self, input_dim, output_dim, hidden_dim = 64):
super(SimpleModel, self).__init__()
self.mlp = MLP(input_dim, output_dim, hidden_dim)
def forward(self, x):
x = self.mlp(x)
return x
Train the model on some data.
\[
y = x_0^2 +3 \sin(x_4)-4
\]
# Make the dataset
x = np.array([np.random.uniform(0, 1, 10_000) for _ in range(5)]).T
y = x[:, 0]**2 + 3*np.sin(x[:, 4]) - 4
noise = np.array([np.random.normal(0, 0.05*np.std(y)) for _ in range(len(y))])
y = y + noise
# Set up training
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.model_selection import train_test_split
def train_model(model, dataloader, opt, criterion, epochs = 100):
"""
Train a model for the specified number of epochs.
Args:
model: PyTorch model to train
dataloader: DataLoader for training data
opt: Optimizer
criterion: Loss function
epochs: Number of training epochs
Returns:
tuple: (trained_model, loss_tracker)
"""
loss_tracker = []
for epoch in range(epochs):
epoch_loss = 0.0
for batch_x, batch_y in dataloader:
# Forward pass
pred = model(batch_x)
loss = criterion(pred, batch_y)
# Backward pass
opt.zero_grad()
loss.backward()
opt.step()
epoch_loss += loss.item()
loss_tracker.append(epoch_loss)
if (epoch + 1) % 5 == 0:
avg_loss = epoch_loss / len(dataloader)
print(f'Epoch [{epoch+1}/{epochs}], Avg Loss: {avg_loss:.6f}')
return model, loss_tracker
# Instantiate the model
model = SimpleModel(input_dim=x.shape[1], output_dim=1)
# Set up training
criterion = nn.MSELoss()
opt = optim.Adam(model.parameters(), lr=0.001)
X_train, _, y_train, _ = train_test_split(x, y.reshape(-1,1), test_size=0.2, random_state=290402)
# Set up dataset
dataset = TensorDataset(torch.FloatTensor(X_train), torch.FloatTensor(y_train))
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# Train the model and save the weights
model, losses = train_model(model, dataloader, opt, criterion, 20)
torch.save(model.state_dict(), 'model_weights.pth')
Epoch [5/20], Avg Loss: 0.076408
Epoch [10/20], Avg Loss: 0.052064
Epoch [15/20], Avg Loss: 0.042181
Epoch [20/20], Avg Loss: 0.033330
Wrap the mlp layer in the trained model with SymbolicMLP.
from symtorch import SymbolicModel
model.mlp = SymbolicModel(model.mlp, block_name = 'Sequential')
Detected IPython. Loading juliacall extension. See https://juliapy.github.io/PythonCall.jl/stable/compat/#IPython
Interpret the MLPο
In this example, we pass extra parameters into the .distill method (complexity of operators/constants and parsimony, which is a penalisation of complexity).
To see all the possible parameters, please see the PySRRegressor class from PySR.
In this example, we turn verbosity off because we are in a Jupyter notebook. For best performance, run in IPython, as you can terminate the SR any time.
# Configure the SR
sr_params = {'complexity_of_operators': {"sin":3, "exp":3},
'complexity_of_constants': 2,
'parsimony': 0.1,
'verbosity': 0,
'niterations': 50}
model.mlp.distill(torch.FloatTensor(X_train),
sr_params = sr_params,
parent_model=model) #Pass in the parent model (really only required if the MLP is
#not at the start of the model but it is good practice)
π οΈ Running SR on output dimension 0 of 0
/Users/liz/PhD/SymTorch_project/symtorch_venv/lib/python3.11/site-packages/pysr/sr.py:2811: UserWarning: Note: it looks like you are running in Jupyter. The progress bar will be turned off.
warnings.warn(
π‘Best equation for output 0 found to be (((exp(x0) * 0.24175334) + -1.879353) + x4) * 2.382997.
β€οΈ SR on Sequential complete.
{0: PySRRegressor.equations_ = [
pick score equation \
0 0.000000e+00 x4
1 2.675247e+00 -2.3036592
2 4.113613e-01 x4 + -2.8012233
3 4.553621e-01 (x4 + x4) + -3.2988408
4 1.302947e-01 (x4 * 2.3804636) + -3.4881692
5 1.391758e+00 x4 + ((x4 + -3.7980285) + x0)
6 8.372327e-01 (x0 + (x4 * 2.3829083)) + -3.988571
7 5.373152e-08 inv(inv((x0 + (x4 * 2.3829083)) + -3.988571))
8 1.165584e-02 ((x4 + (x0 * 0.40451804)) + -1.6661206) * 2.38...
9 1.494943e-06 ((x0 + -3.9705966) + (x4 * 2.3829074)) + (x0 *...
10 >>>> 3.965307e-01 (((exp(x0) * 0.24175334) + -1.879353) + x4) * ...
11 4.405181e-02 (x0 * exp(x0 + -1.0060903)) + ((x4 + -1.617400...
12 9.909887e-04 (((x4 * 0.36639503) + ((x4 + x4) + (x3 * (x4 *...
loss complexity
0 8.091652 1
1 0.557432 2
2 0.244843 4
3 0.098484 6
4 0.086453 7
5 0.021495 8
6 0.009306 9
7 0.009306 11
8 0.009198 12
9 0.009198 14
10 0.006187 15
11 0.005665 17
12 0.005604 28
]}
See the full Pareto front of equations. The best equation is chosen as a balance of accuracy and complexity.
Outputs from PySR are saved in SR_output/MLP_name.
model.mlp.show_symbolic_expression()
β‘οΈ Symbolic expressions for output dimension 0:
complexity loss equation \
0 1 7.717672 x4
1 2 0.482811 -2.2379751
2 4 0.193551 x4 + -2.7430148
3 6 0.069146 (x4 + x4) + -3.2480512
4 7 0.063802 (x4 * 2.2548008) + -3.3767343
5 8 0.014271 x0 + ((x4 + x4) + -3.7485306)
6 9 0.009059 (x4 * 2.2514427) + (x0 + -3.8755205)
7 11 0.006982 (x4 * 2.2504792) + ((x0 * x0) + -3.709683)
8 12 0.006277 (x0 + (sin(x4) * 2.6363063)) + -3.962223
9 14 0.003978 ((x0 * x0) * 0.8169185) + ((x4 * 2.2512612) + ...
10 17 0.001264 ((x0 * sin(x0)) + -3.764245) + (sin(x4) * 2.63...
11 19 0.001157 inv(inv((sin(x4) * 2.6362014) + (((x0 * x0) * ...
12 22 0.000922 inv(sin(inv(((sin(x4) + ((x0 * x0) * 0.3172727...
13 25 0.000879 inv(sin(inv(((sin(x4) * 2.77826) + -3.7736876)...
14 28 0.000873 inv(sin(inv((sin(x4) * 2.7769802) + ((exp(x0 *...
score sympy_format \
0 0.000000 x4
1 2.771643 -2.23797510000000
2 0.457042 x4 - 2.7430148
3 0.514658 x4 + x4 - 3.2480512
4 0.080440 x4*2.2548008 - 3.3767343
5 1.497572 x0 + x4 + x4 - 3.7485306
6 0.454409 x0 + x4*2.2514427 - 3.8755205
7 0.130204 x0*x0 + x4*2.2504792 - 3.709683
8 0.106513 x0 + sin(x4)*2.6363063 - 3.962223
9 0.228033 x0*x0*0.8169185 + x4*2.2512612 - 3.6487224
10 0.382286 x0*sin(x0) + sin(x4)*2.6359825 - 3.764245
11 0.044174 x0*x0*0.81745505 + sin(x4)*2.6362014 - 3.735648
12 0.075511 1/sin(1/((x0*x0*0.31727275 + sin(x4))*2.777514...
13 0.015831 1/sin(1/(exp(x0)*0.3389404*x0 + sin(x4)*2.7782...
14 0.002625 1/sin(1/(exp(x0*1.0814174)*x0*0.3136479 + sin(...
lambda_format
0 PySRFunction(X=>x4)
1 PySRFunction(X=>-2.23797510000000)
2 PySRFunction(X=>x4 - 2.7430148)
3 PySRFunction(X=>x4 + x4 - 3.2480512)
4 PySRFunction(X=>x4*2.2548008 - 3.3767343)
5 PySRFunction(X=>x0 + x4 + x4 - 3.7485306)
6 PySRFunction(X=>x0 + x4*2.2514427 - 3.8755205)
7 PySRFunction(X=>x0*x0 + x4*2.2504792 - 3.709683)
8 PySRFunction(X=>x0 + sin(x4)*2.6363063 - 3.962...
9 PySRFunction(X=>x0*x0*0.8169185 + x4*2.2512612...
10 PySRFunction(X=>x0*sin(x0) + sin(x4)*2.6359825...
11 PySRFunction(X=>x0*x0*0.81745505 + sin(x4)*2.6...
12 PySRFunction(X=>1/sin(1/((x0*x0*0.31727275 + s...
13 PySRFunction(X=>1/sin(1/(exp(x0)*0.3389404*x0 ...
14 PySRFunction(X=>1/sin(1/(exp(x0*1.0814174)*x0*...
π Best: ((x0 * sin(x0)) + -3.764245) + (sin(x4) * 2.6359825) (loss: 1.263592e-03)
Switch to Using the Equation Instead in the Forwards Passο
You can choose the equation you want to switch to by choosing the desired complexity of equation.
If left blank, then we choose the best equation chosen by PySR.
model.mlp.switch_to_symbolic(complexity=[14])
β
Successfully switched Sequential to symbolic equations for all 1 dimensions:
Dimension 0: ((x0 * x0) * 0.8169185) + ((x4 * 2.2512612) + -3.6487224)
Variables: ['x0', 'x4']
π― All 1 output dimensions now using symbolic equations.
Now when running the forwards pass through the model, it uses the symbolic equation instead of the MLP.
interpretable_outputs = model(torch.tensor(X_train, dtype=torch.float32))
interpretable_outputs
tensor([[-1.6968],
[-2.0715],
[-1.1126],
...,
[-2.6302],
[-2.6785],
[-2.9952]])
You can also make a callable function using the get_symbolic_function method.
symbolic_function = model.mlp.get_symbolic_function(complexity = 14)
symbolic_function(X_train)
array([-1.6968284, -2.0715022, -1.1125848, ..., -2.6301584, -2.6784554,
-2.9951892], shape=(8000,), dtype=float32)
Switch to Using the MLP in the Forwards Passο
mlp_outputs = model.mlp.switch_to_block()
β
Switched Sequential back to block
model.eval()
with torch.no_grad():
model_outputs = model.mlp(torch.tensor(X_train, dtype=torch.float32))
model_outputs
tensor([[-1.6186],
[-2.0624],
[-1.1711],
...,
[-2.6357],
[-2.7018],
[-3.0245]])
Wrapping a Python Functionο
SymbolicModel also works with generic Python functions
# Same function as before
def f(X):
return X[:, 0]**2 + 3*np.sin(X[:, 4]) - 4
symbolic_function = SymbolicModel(f, block_name = "callable_function")
symbolic_function.distill(X_train, sr_params=sr_params)
π οΈ Running SR on output dimension 0 of 0
/Users/liz/PhD/SymTorch_project/symtorch_venv/lib/python3.11/site-packages/pysr/sr.py:2811: UserWarning: Note: it looks like you are running in Jupyter. The progress bar will be turned off.
warnings.warn(
π‘Best equation for output 0 found to be (sin(x4) + (x0 * x0)) + ((sin(x4) + sin(x4)) + -4.0).
β€οΈ SR on callable_function complete.
{0: PySRRegressor.equations_ = [
pick score equation \
0 0.000000e+00 x4
1 2.533145e+00 -2.272284
2 3.840348e-01 x4 + -2.777311
3 4.541706e-01 (x4 + x4) + -3.2823517
4 2.530731e-01 (x4 * 2.5681043) + -3.569265
5 9.782338e-01 (x4 + x0) + (x4 + -3.7828333)
6 1.416253e+00 x0 + ((x4 * 2.564849) + -4.068105)
7 5.468644e-01 (x0 * x0) + ((x4 * 2.563873) + -3.9022593)
8 1.771850e-08 ((x4 * 0.56387687) + ((x4 + -3.902262) + (x0 *...
9 3.764534e+00 (((x4 * -0.7083076) + 3.2762177) * x4) + ((x0 ...
10 2.500752e+00 ((inv(x4 + -1.958775) + 3.5588908) * x4) + ((x...
11 6.710256e-07 inv(inv(((inv(-1.958775 + x4) + 3.5588908) * x...
12 >>>> 9.804187e+00 (sin(x4) + (x0 * x0)) + ((sin(x4) + sin(x4)) +...
loss complexity
0 8.008494e+00 1
1 6.359454e-01 2
2 2.950200e-01 4
3 1.189498e-01 6
4 9.235398e-02 7
5 3.472275e-02 8
6 8.424487e-03 9
7 2.821909e-03 11
8 2.821909e-03 15
9 6.540737e-05 16
10 5.364925e-06 17
11 5.364918e-06 19
12 1.635891e-14 21
]}
Note: if distilling a generic function (not a PyTorch MLP) you cannot pass in the parent_model parameter. If this function is part of a hybrid PyTorch model and is not a nn.Module type, then you must pass in the inputs to the function not to the whole model when running .distill.
symbolic_function.switch_to_symbolic()
β
Successfully switched callable_function to symbolic equations for all 1 dimensions:
Dimension 0: (sin(x4) + (x0 * x0)) + ((sin(x4) + sin(x4)) + -4.0)
Variables: ['x0', 'x4']
π― All 1 output dimensions now using symbolic equations.
# Clean up
import os
import shutil
if os.path.exists('SR_output'):
shutil.rmtree('SR_output')
os.remove('model_weights.pth')
if os.path.exists('symtorch_data'):
shutil.rmtree('symtorch_data')