PyTorch Frontend

pyepo.func provides the PyTorch training methods. Each method is an autograd module that wraps an optModel. The forward pass solves the optimization problem. The backward pass applies the method’s gradient rule. Training uses standard PyTorch optimizers.

Method selection and per-method training loops are in Training Methods. This page shows the calling conventions.

Training

End-to-end training of a shortest-path predictor on a 5x5 grid with the SPO+ loss:

import torch
from torch import nn
from torch.utils.data import DataLoader

import pyepo
from pyepo.data.dataset import optDataset

# optimization model: 5x5 grid shortest path
grid = (5, 5)
optmodel = pyepo.model.shortestPathModel(grid)

# synthetic data
x, c = pyepo.data.shortestpath.genData(
    num_data=1000, num_features=5, grid=grid, deg=4, noise_width=0.5, seed=135,
)
dataset = optDataset(optmodel, x, c)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# linear predictor and SPO+ loss
predmodel = nn.Linear(5, optmodel.num_cost)
spo = pyepo.func.SPOPlus(optmodel, processes=1)
optimizer = torch.optim.Adam(predmodel.parameters(), lr=1e-3)

# end-to-end training
for epoch in range(10):
    for xb, cb, wb, zb in dataloader:
        loss = spo(predmodel(xb), cb, wb, zb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

The optDataset batch is (x, c, w, z):

  • x: input features for the prediction model.

  • c: ground-truth objective coefficients.

  • w: optimal solution under c.

  • z: optimal objective value under c.

Some methods use only a subset of these values. The per-method inputs are listed in the summary table of Training Methods.

Solution-Returning Modules

Solution-returning modules such as DPO are trained through a task loss on their output. Perturbed modules draw noise internally. Pass seed= for reproducibility.

dpo = pyepo.func.DPO(optmodel, n_samples=10, sigma=0.5, processes=1)
criterion = nn.MSELoss()

for epoch in range(10):
    for xb, cb, wb, zb in dataloader:
        we = dpo(predmodel(xb))     # expected perturbed solutions
        loss = criterion(we, wb)    # task loss on the solutions
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

GPU

The predictor and the batch can live on CUDA. Every backend except MPAX solves on the CPU. MPAX solves the batch on the GPU (see Solver Backends). The losses expect all tensor inputs on one device, so move the whole batch together. A CPU backend receives CPU copies internally, and the loss and gradients are returned on the batch’s device:

device = "cuda"
predmodel = predmodel.to(device)

for xb, cb, wb, zb in dataloader:
    xb, cb, wb, zb = (t.to(device) for t in (xb, cb, wb, zb))
    loss = spo(predmodel(xb), cb, wb, zb)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Evaluation

pyepo.metric.regret evaluates decision quality, usually on a held-out test set:

total_regret = pyepo.metric.regret(predmodel, optmodel, testloader)

Shared Options

Common constructor options:

  • processes controls the worker pool used for batch solving.

  • solve_ratio enables solution-pool caching when set below 1. CaVE uses this option for its projection branch instead. See CaVE.

  • dataset seeds the solution pool whenever solve_ratio < 1. Contrastive and ranking methods always require it.

  • reduction controls how per-instance losses are aggregated when the method supports it.