PyEPO：基于 PyTorch 的端到端预测-优化工具

学习框架

出版物

本仓库是以下论文的官方实现： PyEPO：用于线性和整数规划的基于 PyTorch 的端到端预测-优化库（已被《数学规划计算》(MPC) 接收）

引用：

@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}
}

简介

PyEPO（基于 PyTorch 的端到端预测-优化工具）是一款基于 Python 的开源软件，支持建模和求解具有线性目标函数的预测-优化问题。PyEPO 的核心功能是利用 GurobiPy、COPT、Pyomo、Google OR-Tools、MPAX 或其他求解器与算法构建优化模型，然后将该优化模型嵌入人工神经网络中进行端到端训练。为此，PyEPO 将多种方法实现为 PyTorch 自动微分模块。

特别地，PyEPO 集成了 MPAX，这是一款基于 JAX 的数学规划求解器，采用 PDHG（原始-对偶混合梯度）算法进行 GPU 加速优化。MPAX 为端到端训练带来了三大优势：(1) GPU 原生求解 — 一阶 PDHG 方法本身具有并行性，可在 GPU 上高效运行；(2) 批量求解 — 整个 mini-batch 的优化实例可通过向量化在 GPU 上同时求解；以及 (3) 无需 GPU–CPU 数据传输开销 — 传统求解器（如 Gurobi）运行在 CPU 上，每次训练迭代都需要在 GPU 和 CPU 之间进行昂贵的数据传输，而 MPAX 则将神经网络和求解器都保留在 GPU 上，从而消除了这一瓶颈。

文档

PyEPO 的官方文档可在 https://khalil-research.github.io/PyEPO 查看。

幻灯片

我们最近的教程是在 ACC 2024 大会上进行的。您可以通过此链接查看演讲幻灯片：https://github.com/khalil-research/PyEPO/blob/main/slides/PyEPO.pdf。

教程

• 01 优化模型： 构建优化求解器
• 02 优化数据集： 生成合成数据并使用 optDataset
• 03 训练与测试： 训练和测试不同的方法
• 04 二维背包问题解可视化： 可视化背包问题的解
• 05 战术争霸最短路径： 在 Warcraft 地形数据集上训练最短路径模型
• 06 实际能源调度： 将 PyEPO 应用于真实能源数据
• 07 kNN 鲁棒损失： 使用 optDatasetKNN 进行鲁棒损失计算
• 08 使用 MPAX 和 PDHG 求解： 使用 MPAX 进行 GPU 加速的批量求解

实验

要复现原论文中的实验，请使用此代码并按照该分支中的说明操作：https://github.com/khalil-research/PyEPO/tree/MPC。请注意，该分支仍处于非常早期的版本阶段。

功能

• 实现 SPO+ [1]、DBB [3]、NID [7]、DPO [4]、PFYL [4]、NCE [5]、LTR [6]、I-MLE [8]、AI-MLE [9] 以及 PG [11]。
• 支持 Gurobi、COPT、Pyomo、Google OR-Tools 和 MPAX 的 API。
• 支持优化求解器的并行计算。
• 支持解缓存 [5]，以加速训练。
• 支持 kNN 稳健损失 [10]，以提高决策质量。

安装

从本仓库克隆并安装

您可以从我们的 GitHub 仓库下载 PyEPO。

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

然后进行安装：

pip install PyEPO/pkg/.

使用 pip 安装

该软件包现已在 PyPI 上架。您可以通过运行以下命令轻松使用 pip 安装 PyEPO：

pip install pyepo

使用 conda 安装

PyEPO 也已在 Anaconda Cloud 上发布。如果您更倾向于使用 conda 进行安装，可以使用以下命令安装 PyEPO：

conda install -c pyepo pyepo

依赖项

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

示例代码

#!/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


# 优化模型
class myModel(optGrbModel):
    def __init__(self, weights):
        self.weights = np.array(weights)
        self.num_item = len(weights[0])
        super().__init__()

    def _getModel(self):
        # 创建模型
        m = gp.Model()
        # 变量
        x = m.addVars(self.num_item, name="x", vtype=GRB.BINARY)
        # 目标函数方向
        m.modelSense = GRB.MAXIMIZE
        # 约束条件
        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


# 预测模型
class LinearRegression(nn.Module):

    def __init__(self):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(num_feat, num_item)

    def forward(self, x):
        out = self.linear(x)
        return out


if __name__ == "__main__":

    # 生成数据
    num_data = 1000 # 数据数量
    num_feat = 5 # 特征维度
    num_item = 10 # 项目数量
    weights, x, c = pyepo.data.knapsack.genData(num_data, num_feat, num_item,
                                                dim=3, deg=4, noise_width=0.5, seed=135)

    # 初始化优化模型
    optmodel = myModel(weights)

    # 初始化预测模型
    predmodel = LinearRegression()
    # 设置优化器
    optimizer = torch.optim.Adam(predmodel.parameters(), lr=1e-2)
    # 初始化 SPO+ 损失函数
    spop = pyepo.func.SPOPlus(optmodel, processes=1)

    # 构建数据集
    dataset = pyepo.data.dataset.optDataset(optmodel, x, c)
    # 获取数据加载器
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

    # 训练
    num_epochs = 10
    for epoch in range(num_epochs):
        for data in dataloader:
            x, c, w, z = data
            # 前向传播
            cp = predmodel(x)
            loss = spop(cp, c, w, z)
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    # 评估
    regret = pyepo.metric.regret(predmodel, optmodel, dataloader)
    print("训练集上的遗憾值：{:.4f}".format(regret))

参考文献

• [1] Elmachtoub, A. N., & Grigas, P. (2021). 智能“先预测、后优化”。管理科学。
• [2] Mandi, J., Stuckey, P. J., & Guns, T. (2020). 面对困难组合优化问题的智能预测-优化方法。收录于AAAI人工智能会议论文集。
• [3] Vlastelica, M., Paulus, A., Musil, V., Martius, G., & Rolínek, M. (2019). 黑盒组合求解器的差异化。arXiv预印本，arXiv:1912.02175。
• [4] Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J. P., & Bach, F. (2020). 基于可微扰动优化器的学习。神经信息处理系统进展，第33卷，页9508–9519。
• [5] Mulamba, M., Mandi, J., Diligenti, M., Lombardi, M., Bucarey, V., & Guns, T. (2021). 用于预测-优化的对比损失与解缓存技术。第三十届国际人工智能联合会议论文集。
• [6] Mandi, J., Bucarey, V., Mulamba, M., & Guns, T. (2022). 以决策为中心的学习：从排序学习的视角看。第三十九届国际机器学习会议论文集。
• [7] Sahoo, S. S., Paulus, A., Vlastelica, M., Musil, V., Kuleshov, V., & Martius, G. (2022). 通过组合算法进行反向传播：带投影的恒等映射有效。arXiv预印本，arXiv:2205.15213。
• [8] Niepert, M., Minervini, P., & Franceschi, L. (2021). 隐式最大似然估计：通过离散指数族分布进行反向传播。神经信息处理系统进展，第34卷，页14567–14579。
• [9] Minervini, P., Franceschi, L., & Niepert, M. (2023年6月). 面向离散潜在变量模型的自适应扰动梯度估计。收录于AAAI人工智能会议论文集（第37卷，第8期，页9200–9208）。
• [10] Schutte, N., Postek, K., & Yorke-Smith, N. (2023). 用于以决策为中心学习的鲁棒损失函数。arXiv预印本，arXiv:2310.04328。
• [11] Gupta, V., & Huang, M. (2024). 基于方向梯度的以决策为中心的学习。培训，第50卷第100期，第150页。

PyEPO 快速上手指南

PyEPO 是一个基于 PyTorch 的端到端“预测后优化”（Predict-then-Optimize）工具库。它允许用户将线性或整数规划模型嵌入神经网络中，利用 Gurobi、COPT、OR-Tools 或 MPAX 等求解器进行端到端训练，特别支持 GPU 加速的批量求解。

环境准备

系统要求

• 操作系统：Linux, macOS, Windows
• Python 版本：3.8 及以上

前置依赖

在使用 PyEPO 前，请确保已安装以下核心依赖：

• PyTorch：深度学习框架
• 求解器（至少安装其一）：
  • Gurobi (需配置许可证)
  • COPT (杉数科技)
  • Google OR-Tools
  • MPAX (基于 JAX，支持 GPU 加速)
• 其他 Python 库：NumPy, SciPy, Scikit-Learn, Pyomo, Pathos, tqdm

提示：若使用 Gurobi 或 COPT 商业求解器，请提前申请并配置好相应的 License。

安装步骤

你可以通过 pip 或 conda 直接安装稳定版，也可以从源码安装最新开发版。

方式一：使用 pip 安装（推荐）

pip install pyepo

国内用户可使用清华源加速：

pip install pyepo -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：使用 conda 安装

conda install -c pyepo pyepo

方式三：从源码安装

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

基本使用

以下示例演示了如何构建一个简单的背包问题优化模型，并结合线性回归神经网络，使用 SPO+ 损失函数进行端到端训练。

1. 导入依赖与定义模型

#!/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

# 定义优化模型 (继承 optGrbModel)
class myModel(optGrbModel):
    def __init__(self, weights):
        self.weights = np.array(weights)
        self.num_item = len(weights[0])
        super().__init__()

    def _getModel(self):
        m = gp.Model()
        # 定义变量
        x = m.addVars(self.num_item, name="x", vtype=GRB.BINARY)
        # 设置最大化目标
        m.modelSense = GRB.MAXIMIZE
        # 添加约束
        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

# 定义预测模型 (简单的线性回归)
class LinearRegression(nn.Module):
    def __init__(self, num_feat, num_item):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(num_feat, num_item)

    def forward(self, x):
        return self.linear(x)

2. 数据生成与初始化

if __name__ == "__main__":
    # 参数设置
    num_data = 1000
    num_feat = 5
    num_item = 10
    
    # 生成合成数据 (背包问题)
    weights, x, c = pyepo.data.knapsack.genData(num_data, num_feat, num_item,
                                                dim=3, deg=4, noise_width=0.5, seed=135)

    # 初始化优化模型和预测模型
    optmodel = myModel(weights)
    predmodel = LinearRegression(num_feat, num_item)
    
    # 设置优化器
    optimizer = torch.optim.Adam(predmodel.parameters(), lr=1e-2)
    
    # 初始化 SPO+ 损失函数
    spop = pyepo.func.SPOPlus(optmodel, processes=1)

    # 构建数据集和数据加载器
    dataset = pyepo.data.dataset.optDataset(optmodel, x, c)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

3. 训练与评估

    # 训练循环
    num_epochs = 10
    for epoch in range(num_epochs):
        for data in dataloader:
            x, c, w, z = data
            # 前向传播
            cp = predmodel(x)
            loss = spop(cp, c, w, z)
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    # 评估遗憾值 (Regret)
    regret = pyepo.metric.regret(predmodel, optmodel, dataloader)
    print("Regret on Training Set: {:.4f}".format(regret))

运行上述代码即可完成一次完整的端到端训练流程。更多高级用法（如使用 MPAX GPU 加速、不同损失函数对比等）请参考官方文档或 Colab 教程。