Datasets

The optDataset family turns features and costs into a PyTorch Dataset with the optimization labels the training methods need. Each dataset builds those labels by solving an optModel. The model can be a built-in one or one you define yourself.

optDataset

pyepo.data.optDataset is a PyTorch Dataset that stores features and cost coefficients, and solves the optimization problem to obtain optimal solutions and objective values. The features and costs can come from any arrays. The generators in Data Generators are only a convenience.

optDataset precomputes \(\mathbf{w}^*(\mathbf{c})\) and \(z^*(\mathbf{c})\) at construction time. If those labels already exist from another source, you can skip optDataset and feed (x, c, w, z) batches to pyepo.func modules directly.

class pyepo.data.dataset.optDataset(model: optModel, feats: ndarray | Tensor, costs: ndarray | Tensor)

PyTorch Dataset for predict-then-optimize problems.

At construction time it solves the optimization problem for every cost vector and caches the optimal solution \(\mathbf{w}^*(\mathbf{c})\) and objective value \(z^*(\mathbf{c})\). This precomputation removes solver overhead from the training loop, making optDataset the standard input format for end-to-end training in PyEPO. When labels are already available from another source, optDataset can be skipped and batches fed directly to pyepo.func modules.

Variables:

The following example shows how to use optDataset with a PyTorch DataLoader:

import pyepo
from torch.utils.data import DataLoader

# model for shortest path
grid = (5,5) # grid size
model = pyepo.model.shortestPathModel(grid)

# generate data
num_data = 1000 # number of samples
num_feat = 5 # number of features
deg = 4 # polynomial degree
noise_width = 0 # noise width
x, c = pyepo.data.shortestpath.genData(num_data, num_feat, grid, deg, noise_width, seed=135)

# build dataset
dataset = pyepo.data.dataset.optDataset(model, x, c)

# get data loader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

optDatasetKNN

pyepo.data.optDatasetKNN is a PyTorch Dataset that implements the k-nearest neighbors (kNN) robust loss [1] for predict-then-optimize training. It stores features and cost coefficients, then computes the mean k-nearest-neighbor solutions and the corresponding optimal objective values.

For a runnable walkthrough, see the 08 kNN Robust Losses notebook.

class pyepo.data.dataset.optDatasetKNN(model: optModel, feats: ndarray | Tensor, costs: ndarray | Tensor, k: int = 10, weight: float = 0.5)

PyTorch Dataset for the kNN-robust decision-focused loss.

For each instance the cost vector is replaced with a convex combination of its k nearest neighbours in feature space, and the optimization problem is solved on the smoothed costs. The mean kNN solutions and objective values are cached for training, providing a robust supervision signal under noisy or out-of-distribution feature observations.

Reference: Schutte et al. (2024) https://www.ijcai.org/proceedings/2024/538

Variables:

  • model (optModel) – Optimization model

  • k (int) – number of nearest neighbours selected

  • weight (float) – self-weight in the kNN convex combination (1.0 = no smoothing)

  • feats (torch.Tensor) – Data features

  • costs (torch.Tensor) – kNN-smoothed cost vectors

  • sols (torch.Tensor) – Mean kNN optimal solutions

  • objs (torch.Tensor) – Mean kNN optimal objective values

import pyepo
from torch.utils.data import DataLoader

# model for shortest path
grid = (5,5) # grid size
model = pyepo.model.shortestPathModel(grid)

# generate data
num_data = 1000 # number of samples
num_feat = 5 # number of features
deg = 4 # polynomial degree
noise_width = 0 # noise width
x, c = pyepo.data.shortestpath.genData(num_data, num_feat, grid, deg, noise_width, seed=135)

# build dataset
dataset = pyepo.data.dataset.optDatasetKNN(model, x, c, k=10, weight=0.5)

# get data loader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

optDatasetConstrs

pyepo.data.dataset.optDatasetConstrs is a PyTorch Dataset for the CaVE [2] cone-aligned loss. In addition to the features, costs, optimal solutions, and objective values stored by optDataset, it also extracts the normals of the binding constraints at the optimal vertex for each instance. CaVE then projects the sense-flipped predicted cost vector onto the cone spanned by these normals.

optDatasetConstrs currently requires a Gurobi-backed optModel. The dataset also checks that the optimal vertex is binary, since CaVE is defined for binary linear programs. It raises on infeasible instances or non-binary optima unless skip_infeas=True, which drops them instead.

For a runnable walkthrough that uses optDatasetConstrs end-to-end with the CaVE loss, see the 04 CaVE for Binary Linear Programs notebook.

class pyepo.data.dataset.optDatasetConstrs(model: optModel, feats: ndarray | Tensor, costs: ndarray | Tensor, skip_infeas: bool = False)

PyTorch Dataset for the CaVE cone-aligned loss.

Stores features and cost coefficients, solves each instance, and extracts the normals of the binding constraints at the optimal vertex in canonical <= orientation. These normals span the polyhedral cone onto which coneAlignedCosine projects the predicted cost vector during training.

CaVE is defined for binary linear programs only, so the optimal vertex must be binary; instances that are infeasible or have non-binary optima raise (or are skipped when skip_infeas=True). Binding-constraint extraction uses Gurobi’s sparse-matrix API, which is why this dataset currently requires a Gurobi-backed optModel.

Per-instance row counts differ (different constraints bind at different vertices), so batch with optDataLoader or pass collate_tight_constraints to a PyTorch DataLoader.

Reference: Tang & Khalil (2024) https://link.springer.com/chapter/10.1007/978-3-031-60599-4_12

Variables:

Per-instance constraint matrices can have different row counts because different constraints bind at different vertices. Batch them with optDataLoader, which pads them automatically:

class pyepo.data.dataset.optDataLoader(dataset, *args, **kwargs)

DataLoader that applies a dataset’s own collate_fn when present.

Datasets with ragged samples (e.g. optDatasetConstrs, whose binding-constraint matrices vary in row count) carry a collate_fn; this loader uses it so the caller never passes one explicitly. Plain datasets fall back to the default PyTorch collation.

import pyepo
from pyepo.data.dataset import optDatasetConstrs, optDataLoader

# model for TSP (Gurobi backend required)
model = pyepo.model.tspModel(num_nodes=10, formulation="DFJ")

# generate data
x, c = pyepo.data.tsp.genData(num_data=1000, num_features=5, num_nodes=10, deg=4, seed=135)

# build CaVE dataset (extracts tight binding-constraint normals at the optimum)
dataset_constr = optDatasetConstrs(model, x, c)

# optDataLoader pads ragged per-instance constraint matrices
dataloader_constr = optDataLoader(dataset_constr, batch_size=32, shuffle=True)