PyEPO: A PyTorch-based End-to-End Predict-then-Optimize Tool

<p align="center"><img width="100%" src="images/logo1.png" /></p>

Learning Framework

<p align="center"><img width="100%" src="images/learning_framework_e2e.png" /></p>

Publication

This repository is the official implementation of the paper: PyEPO: A PyTorch-based End-to-End Predict-then-Optimize Library for Linear and Integer Programming (Accepted to Mathematical Programming Computation (MPC))

Citation:

@article{tang2024,
  title={PyEPO: a PyTorch-based end-to-end predict-then-optimize library for linear and integer programming},
  author={Tang, Bo and Khalil, Elias B},
  journal={Mathematical Programming Computation},
  issn={1867-2957},
  doi={10.1007/s12532-024-00255-x},
  year={2024},
  month={July},
  publisher={Springer}
}

Introduction

PyEPO (PyTorch-based End-to-End Predict-then-Optimize Tool) is a Python-based, open-source software that supports modeling and solving predict-then-optimize problems with linear objective functions. The core capability of PyEPO is to build optimization models with GurobiPy, COPT, Pyomo, Google OR-Tools, MPAX or any other solvers and algorithms, then embed the optimization model into an artificial neural network for the end-to-end training. For this purpose, PyEPO implements various methods as PyTorch autograd modules.

In particular, PyEPO integrates MPAX, a JAX-based mathematical programming solver using the PDHG (Primal-Dual Hybrid Gradient) algorithm for GPU-accelerated optimization. MPAX brings three key advantages for end-to-end training: (1) GPU-native solving — the first-order PDHG method is inherently parallelizable and runs efficiently on GPU; (2) batch solving — an entire mini-batch of optimization instances can be solved simultaneously on GPU via vectorization; and (3) no GPU–CPU data transfer overhead — traditional solvers (e.g., Gurobi) run on CPU, requiring costly data transfers between GPU and CPU at every training iteration, whereas MPAX keeps both the neural network and the solver on GPU, eliminating this bottleneck.

Documentation

The official PyEPO docs can be found at https://khalil-research.github.io/PyEPO.

Slides

Our recent tutorial was at the ACC 2024 conference. You can view the talk slides here.

Tutorial

• 01 Optimization Model: Build an optimization solver
• 02 Optimization Dataset: Generate synthetic data and use optDataset
• 03 Training and Testing: Train and test different approaches
• 04 2D knapsack Solution Visualization: Visualize solutions for the knapsack problem
• 05 Warcraft Shortest Path: Train shortest path models on the Warcraft terrains dataset
• 06 Real-World Energy Scheduling: Apply PyEPO to real energy data
• 07 kNN Robust Losses: Use optDatasetKNN for robust losses
• 08 Solving on MPAX with PDHG: Use MPAX for GPU-accelerated batch solving

Experiments

To reproduce the experiments in the original paper, please use the code and follow the instructions in this branch. Please note that this branch is a very early version.

Features

• Implement SPO+ [1], DBB [3], NID [7], DPO [4], PFYL [4], NCE [5], LTR [6], I-MLE [8], AI-MLE [9], and PG [11].
• Support Gurobi, COPT, Pyomo, Google OR-Tools and MPAX API
• Support parallel computing for optimization solvers
• Support solution caching [5] to speed up training
• Support kNN robust loss [10] to improve decision quality

Installation

Clone and Install from this Repo

You can download PyEPO from our GitHub repository.

git clone -b main --depth 1 https://github.com/khalil-research/PyEPO.git

And install it.

pip install PyEPO/pkg/.

Pip Install

The package is now available for installation on PyPI. You can easily install PyEPO using pip by running the following command:

pip install pyepo

Conda Install

PyEPO is also available on Anaconda Cloud. If you prefer to use conda for installation, you can install PyEPO with the following command:

conda install -c pyepo pyepo

Dependencies

• NumPy
• SciPy
• Pathos
• tqdm
• Pyomo
• Gurobi
• Scikit Learn
• PyTorch

Sample Code

#!/usr/bin/env python
# coding: utf-8

import gurobipy as gp
from gurobipy import GRB
import numpy as np
import pyepo
from pyepo.model.grb import optGrbModel
import torch
from torch import nn
from torch.utils.data import DataLoader


# optimization model
class myModel(optGrbModel):
    def __init__(self, weights):
        self.weights = np.array(weights)
        self.num_item = len(weights[0])
        super().__init__()

    def _getModel(self):
        # create a model
        m = gp.Model()
        # variables
        x = m.addVars(self.num_item, name="x", vtype=GRB.BINARY)
        # model sense
        m.modelSense = GRB.MAXIMIZE
        # constraints
        m.addConstr(gp.quicksum([self.weights[0,i] * x[i] for i in range(self.num_item)]) <= 7)
        m.addConstr(gp.quicksum([self.weights[1,i] * x[i] for i in range(self.num_item)]) <= 8)
        m.addConstr(gp.quicksum([self.weights[2,i] * x[i] for i in range(self.num_item)]) <= 9)
        return m, x


# prediction model
class LinearRegression(nn.Module):

    def __init__(self):
        super(LinearRegression, self).__