ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

Windows 安装跑通 BFM 与 Wasserstein 梯度流实战

Windows 安装跑通 BFM 与 Wasserstein 梯度流实战

把一套在 Linux 服务器上跑得很顺的Wasserstein 梯度流数值实验代码搬到 Windows 工作站上,是我折腾时间最久的一件小事。算法本身并不复杂,核心就是 JKO 格式加一个最优传输求解器,而BFM(back-and-forth method)恰好是这里最巧的一环——它用一个前后交替的迭代,把每一步昂贵的传输映射求解拆成一系列加权最近邻查询,绕开了网格离散带来的维度灾难。麻烦的地方在于,标题里写的“windows安装”才是真正的拦路虎:数学推导看两小时能懂,环境配置能卡三天。这篇内容就是把我把The Back-And-Forth Method (BFM) for Wasserstein Gradient Flows这套东西在 Windows 上从零跑通的全过程摊开讲,包括方案选型的取舍、可直接抄的命令、一维和二维两套验证代码、参数怎么定,以及我踩过的那些报错和数值异常。适合已经懂一点最优传输、想在本地快速复现实验的算法工程师和研究生,也适合只想知道“这套东西值不值得在 Windows 上搭”的旁观者。

1. 先搞清楚要装的是什么:BFM 与 Wasserstein 梯度流的真实关系

在动手conda create之前,我建议先花二十分钟把这三者的关系理顺:Wasserstein 梯度流是“要算的东西”,JKO 是“怎么把它离散成一步步的”,BFM 是“每一步的那个最优传输问题怎么快速解”。这三个概念混在一起,是后面调试时最容易被误导的根源——比如能量不下降,有人会怀疑 BFM 写错了,实际上大概率是时间步 τ 取大了,跟传输求解器没关系。我在第一次复现时就干过这种蠢事,把 JKO 的迭代次数调了三遍,最后发现是势能梯度写错了符号。

1.1 Wasserstein 梯度流到底在算什么

把一堆粒子想象成沙子,Wasserstein 梯度流描述的就是“沙堆在某个能量地形里怎么自然塌陷”。形式化一点说,给定能量泛函 F(ρ),它可能是外势能 ∫V dρ、内部能量 ∫ρ log ρ,或者相互作用能,我们想让概率密度 ρ 沿着让 F 下降最快的方向演化,而“最快”这个距离度量用的是 W2(二次 Wasserstein 距离)。这跟普通梯度下降的区别在于:普通梯度下降在欧氏空间里走直线,Wasserstein 梯度流是在“搬土方最省力”的度量下走,所以它会自然保持质量守恒、不产生负密度、也不会出现密度被“拉断”的假象。这也是为什么做采样、做生成模型、做密度演化的人特别偏爱它——数值行为稳,物理意义清楚。

实际计算的时候,我们没法直接在连续密度上演化,最常见的两种离散方式是网格法和粒子法。网格法在二维还能勉强撑住,三维往上就会被“维数灾难”按在地上摩擦:想算准传输映射,网格点数得指数级增长。粒子法则把密度表示成 N 个带权重的点,传输映射只在点集之间求解,维度带来的额外开销主要体现在最近邻搜索上,而不是网格数量上。BFM 属于后一条路线,而且它特别适合粒子法,因为它输出的直接就是映射,不是一个需要再转换的对偶势。

1.2 BFM 的关键一步:用前后交替代替网格求解

最优传输的核心是 Brenier 定理:在二次代价下,最优映射是某个凸函数的梯度。麻烦在于这个凸函数怎么求。经典做法是解 Monge-Ampère 方程,这玩意是高度非线性的,离散后还得处理凸性约束,工程上非常难受。另一条路是 Kantorovich 的对偶,转成一个线性规划或者半离散的凸优化,好处是凸,坏处是变量数一上去就爆。

BFM 的思路很“工程师”:既然前向映射和后向映射互为某种逆,那就不如交替地求。具体地说,固定一个方向的对偶势,另一个方向的映射就能写成一个加权的最近邻问题——对每个源点 x_i,找让 |x_i − y_j|²/2 − g_j 最小的那个目标点 y_j,这就是所谓的 c-transform,也是半离散最优传输里“功率图”的标准写法。反过来再固定前向势去更新后向映射,如此来回。每轮迭代只需要做一次最近邻查询,用 k-d 树就是 O(N log N),没有网格,没有线性规划,内存也是可控的。我自己的理解是:BFM 把“解一个全局的传输问题”换成了“反复做便宜但不够准的局部查询”,用迭代次数换取单次成本,这在粒子数几千到几万的规模上非常划算。

这套东西用在梯度流上的自然方式,就是嵌进 JKO(minimizing movement)格式:每一步先解一个从当前分布到新分布的传输问题,再在传输后的位置上做一次能量下降。BFM 在这个循环里充当“传输求解器”的角色。需要说明的是,原始论文里的格式更严谨,包含了映射与能量之间的交替最小化以及相应的收敛性分析;我在下面给出的实现是在这个框架下做的工程化简化版,目的是让环境验证和手感建立能在几分钟内完成,严格的数值格式和收敛阶还是以原文为准。

1.3 为什么“安装”这件事在 Windows 上会变复杂

如果在 Linux 上,这套东西的安装基本是三行命令:建环境、装 numpy/scipy、装个 POT 做对照,完事。Windows 上麻烦的地方在于好几个点叠在一起。

第一是二进制轮子的可用性。numpy 和 scipy 现在官方都提供 Windows 的预编译 wheel,直接 pip 装就行;但一旦碰到需要编译的包(比如某些最优传输的精确求解器、或者带 C++ 扩展的加速库),就得有 MSVC 的 C++ 生成工具链,而 Linux 上只要 gcc 就在那儿躺着。第二是并行模型差异:Linux 用 fork,Windows 用 spawn,多进程代码在 Linux 上跑得好好的,到 Windows 上会莫名其妙地把整个模块重新导入一遍,如果没有if __name__ == "__main__":保护,直接无限递归创建进程,我第一次见到这个报错时盯着屏幕愣了半分钟。第三是路径与编码:中文用户名、带空格的目录、OneDrive 同步目录、cp936 默认编码,每一条都能让 matplotlib 或者文件读取出问题。

还有一条容易被忽略:Windows Defender 的实时扫描会显著拖慢第一次 import numpy 的加载速度,因为它在扫描那些几十兆的 DLL。这个不是错误,但会让你误以为环境坏了。我的处理方式是把 conda 环境目录和项目目录加进排除列表,加载时间从十几秒降到两秒出头。

2. 环境方案设计:Windows 上跑 BFM 的最小可用栈

方案选型这件事我走过弯路,所以先把结论放前面:用 conda 系发行版建独立环境,Python 锁在 3.10~3.11,numpy/scipy/matplotlib 走 conda 安装,纯 Python 包(比如 POT)走 pip,编译工具链按需再装。这个组合在我试过的几台机器上都是最省事的。下面把三种常见路线的取舍摊开讲,你可以对号入座。

2.1 三条路线的取舍:裸 pip、conda、WSL

维度裸 Python + venvconda 系发行版WSL2
上手成本低,但要自己管 Python 版本中,多一层概念高,要配发行版和文件系统
二进制依赖靠 PyPI wheel,缺了就编译有 conda-forge 预编译包等同于 Linux,最省心
BLAS 后端跟 wheel 绑定,不可控可选 MKL / OpenBLAS同 Linux
多进程行为spawn,需显式保护spawn,需显式保护fork,老代码直接能跑
图形显示原生,最顺原生,最顺需 X server,画图麻烦
适合场景只跑纯 Python 小脚本主力推荐需要复现 Linux 环境

我最后选 conda 的核心原因是 BLAS 可控。Wasserstein 梯度流的粒子格式里,最热的循环是距离计算和最近邻查询,这些都是 BLAS 和底层 C 代码在跑。裸 pip 装的 numpy 在 Windows 上用的是一套打包好的 OpenBLAS,理论上也能用,但一旦和 MKL 版本的其他库混在一起,就容易出现重复的 OpenMP 运行时,报那个经典的OMP: Error #15: Initializing libiomp5md.dll。conda 帮你把所有包的 BLAS 后端统一到一套,这类冲突基本消失。

WSL2 我没有作为主力,原因很朴素:我要画收敛曲线、要看粒子分布,图形界面在 WSL 里要额外折腾,而且跨文件系统读写大数组的性能有明显损耗。如果你只是跑批处理、把结果存成 npy 再拿到 Windows 上看,WSL2 是很香的,尤其是复现那些只在 Linux 上测试过的老仓库时。

2.2 一次建好、长期复用的虚拟环境

具体的建立流程我整理成了一套可以反复用的模板,注意路径要选纯英文、不带空格、不在 OneDrive 同步目录下——同步盘会锁文件,conda 装到一半报“另一个程序正在使用此文件”,排查起来非常浪费时间。

# 1. 建环境,Python 版本别用最新的,3.10/3.11 生态最稳 conda create -n bfm python=3.11 -y # 2. 激活(PowerShell 下如果报执行策略错误,见 4.1 的排查表) conda activate bfm # 3. 科学计算三件套走 conda,一次装齐省得后面版本打架 conda install -c conda-forge numpy scipy matplotlib -y # 4. 需要对照精确解时再装 POT(纯 Python + 可选扩展,pip 装即可) pip install POT # 5. 验证 python -c "import numpy, scipy, matplotlib; print(numpy.__version__, scipy.__version__)"

这几条命令里有几个点值得单独说。-c conda-forge建议显式写上,defaults 频道的包版本往往落后,而且 conda-forge 的 Windows 构建质量这两年明显更好。pip install POT不要写成pip install ot,后者是另一个完全不相干的包,装上了 import 会给你一个惊喜。装完之后一定要跑一次python -c "import ..."验证,别等写完代码再发现 import 失败。

环境建好之后,我习惯再导出一份environment.yml放进项目根目录。这不是形式主义,是因为这类数值实验经常要过几个月再回来重跑,那时候包版本早就变了,numpy 2.x 把np.float_、np.alltrue这些别名删掉,一堆老代码直接崩,你会感谢当时留了一份快照。

2.3 什么时候真的需要 C++ 编译工具链

大部分情况下你不需要。numpy、scipy、matplotlib、POT 在 Windows 上都有现成的 wheel,pip install就是下载解压。真正需要编译工具的典型场景是:某个包没有对应 Python 版本的 wheel、你想从源码装一个开发版、或者你在写 Cython / pybind11 扩展做加速。

判断方法很简单:pip install的时候如果看到Building wheel for xxx并且开始刷cl.exe相关的报错,那就是缺工具链。这时装一份 Visual Studio Build Tools,勾选“使用 C++ 的桌面开发”即可,注意不需要装完整的 IDE。装完之后在新开的终端里跑where cl确认能找到编译器。

注意:conda 环境里如果同时存在 conda 装的编译器和系统装的 MSVC,可能出现链接到错误运行时库的问题。我的一般做法是二选一,优先用系统 MSVC,不额外装 conda 的编译器元包。

还有一个小坑:如果你用的是较新的 Python 版本,某些老包的 setup.py 会调distutils,而这个模块在新版本里已经被移除,报ModuleNotFoundError: No module named 'distutils'。解决办法是装setuptools(通常已经装了),或者干脆换一个提供 wheel 的包版本,别在这上面死磕。

3. 分步实操:从空机器到第一个收敛曲线

环境搭好只是开始,真正验证这套东西“装对了”的标准,是能跑出一条能量单调下降的曲线,并且和解析解对得上。我设计了两级验证:一维用有闭式解的二次势能锁死正确性,二维用点云做传输映射的对照测试。这种做法在数值实验里很常见——先用有解析解的最简情形把 bug 挤出来,再上复杂的。

3.1 从零到 import 成功的十条命令

把前面零散的命令整理成一条完整的时间线,方便你从头走一遍。每一步后面都跟了验证动作,出错能立刻定位到哪一环。

# 1) 确认 conda 可用 conda --version # 2) 建环境 conda create -n bfm python=3.11 -y # 3) 激活 conda activate bfm # 4) 确认 python 指向环境内的解释器,而不是系统那个 where python # 5) 装科学计算栈 conda install -c conda-forge numpy scipy matplotlib -y # 6) 装对照库 pip install POT # 7) 逐项验证 python -c "import numpy as np; print(np.__version__)" python -c "import scipy; print(scipy.__version__)" python -c "from scipy.spatial import cKDTree; print('kdtree ok')" python -c "import ot; print('POT ok')" # 8) 建项目目录(纯英文路径) mkdir C:\bfm_lab cd C:\bfm_lab

第 4 步值得强调。Windows 上很容易出现“激活了环境但 python 还是系统的”这种情况,尤其是装了多个发行版之后。where python输出的第一行应该是环境目录下的路径,比如C:\Users\xxx\miniforge3\envs\bfm\python.exe。如果不是,说明 PATH 顺序有问题,或者激活脚本没生效。

3.2 一维基准:用解析解给环境“体检”

一维是最好的起点,因为二次代价下的最优传输映射就是分位数映射,排序即可精确求解,不存在任何近似。更妙的是,对二次势能 V(x)=x²/2,JKO 格式有闭式解,可以拿来当“标准答案”。

推导很简单:JKO 的一步是求 y 使得 (y−x)²/(2τ) + y²/2 最小,对 y 求导得 (y−x)/τ + y = 0,解出 y = x/(1+τ)。也就是说经过 n 步后 x_n = x_0/(1+τ)^n。而连续流 dx/dt = −x 的解是 x_0 e^{−t},两者在 τ→0 时一阶一致。这给了我们一个非常干净的验证点:数值解应该精确复现 (1+τ)^(−n),并且随 τ 减小逼近指数解。

import numpy as np from scipy.optimize import minimize def jko_1d(x0, w, V, gradV, tau, steps): """一维 JKO 格式:每步解一个逐点的凸优化。 一维下最优传输映射保序,所以无约束求解不会破坏结构。""" x = x0.astype(float).copy() energies = [] for _ in range(steps): f = lambda y: np.sum(w * (y - x) ** 2) / (2 * tau) + np.sum(w * V(y)) g = lambda y: w * (y - x) / tau + w * gradV(y) res = minimize(f, x, jac=g, method="L-BFGS-B", options={"maxiter": 500, "ftol": 1e-14}) x = res.x energies.append(f(x)) return x, np.array(energies) # 案例一:二次势,与闭式解对照 tau, steps = 0.05, 100 x0 = np.array([1.0, -2.0, 0.5]) w = np.ones(3) / 3 V = lambda x: 0.5 * x ** 2 gradV = lambda x: x x_num, E = jko_1d(x0, w, V, gradV, tau, steps) x_ref = x0 / (1 + tau) ** steps print("数值解:", x_num) print("闭式解:", x_ref) print("最大误差:", np.abs(x_num - x_ref).max())

跑出来误差应该在 1e-8 量级。如果误差是 1e-2 或者根本没收敛,八成是对w的处理出了问题——w是每个粒子的权重,求和必须为 1。我第一次写的时候忘了归一化,结果相当于把能量整体放大,L-BFGS-B 在梯度尺度不对的情况下会提前退出,报ABNORMAL_TERMINATION_IN_LNSRCH,看着像算法问题,其实是归一化问题。

案例二换成双井势 V(x) = (x²−1)²/4,梯度是 x³−x。这个势能有两个极小点 ±1,中间有个势垒,是展示“相分离”的经典例子。取 N=4000 个粒子从 N(0, 0.3²) 采样,跑 200 步 JKO,你会看到粒子分成两团,各自往 ±1 靠,中间逐渐空出来。因为一维最优映射保序,粒子不会跨越原点,所以两侧的粒子数量从头到尾不变——这个现象可以拿来检查你的实现有没有破坏保序性,如果你的模拟里出现了粒子穿到对面去,那一定是求解器或者更新方式写错了。

rng = np.random.default_rng(0) N = 4000 x0 = rng.normal(0, 0.3, N) w = np.ones(N) / N V = lambda x: 0.25 * (x ** 2 - 1) ** 2 gradV = lambda x: x ** 3 - x x_final, E = jko_1d(x0, w, V, gradV, tau=0.02, steps=200) print("左团粒子数:", (x_final < 0).sum()) print("右团粒子数:", (x_final > 0).sum()) print("能量是否单调下降:", np.all(np.diff(E) <= 1e-10))

最后那个单调性检查是我每次都会加的断言。JKO 格式在理论上保证能量单调不增,如果你的数值解出现了能量上升,说明离散化破坏了结构,通常是 τ 太大或者求解器没收敛。这条断言比看曲线图靠谱得多。

3.3 二维点云:把 BFM 用起来

一维验证通过,说明环境本身没问题。接下来上二维,这才是 BFM 真正派上用场的地方——一维下排序就能解传输,根本不需要 BFM,二维开始才有意思。

BFM 的核心是前后交替的 c-transform 迭代。我把它写成一个独立的函数,输入两个点云和权重,输出映射索引和对偶势:

import numpy as np def pairwise_sq(X, Y): """分块计算平方距离矩阵,避免 N*M 直接爆内存。""" return ((X[:, None, :] - Y[None, :, :]) ** 2).sum(-1) def ctransform_fwd(X, Y, g): """前向 c-transform:对每个源点找最优目标点。""" D = pairwise_sq(X, Y) vals = 0.5 * D - g[None, :] j = vals.argmin(axis=1) return 0.5 * D[np.arange(len(X)), j] - g[j], j def ctransform_bwd(X, Y, f): """后向 c-transform:对每个目标点找最优源点。""" D = pairwise_sq(X, Y) vals = 0.5 * D - f[:, None] i = vals.argmin(axis=0) return 0.5 * D[i, np.arange(len(Y))] - f[i], i def bfm_map(X, Y, iters=200, theta=0.5): """简化版 BFM:交替更新前向/后向对偶势,阻尼系数 theta 控稳定性。""" f = np.zeros(len(X)) g = np.zeros(len(Y)) for _ in range(iters): f_new, j = ctransform_fwd(X, Y, g) g_new, i = ctransform_bwd(X, Y, f) f += theta * (f_new - f) g += theta * (g_new - g) return j, f, g

几个实现细节值得说清楚。theta是阻尼系数,取 1.0 是纯 c-transform 迭代,收敛但容易来回震荡;取 0.5 相当于对更新做半步平均,在我试过的点云上更稳。iters不用设太大,通常 100~300 轮就能让映射稳定下来,判断标准是前向映射的索引连续若干轮不再变化。距离矩阵那块,N=M=3000 时是 9×10⁶ 个浮点数,约 72MB,还能接受;再往上就得换 cKDTree 做近邻查询,或者分块处理。

验证映射对不对,最直接的办法是拿 POT 的精确解做对照:

import ot import numpy as np rng = np.random.default_rng(1) N = 800 X = rng.normal(0, 1, (N, 2)) Y = rng.normal(1.5, 0.8, (N, 2)) a = np.ones(N) / N b = np.ones(N) / N # 精确解 M = pairwise_sq(X, Y) G = ot.emd(a, b, M) j_exact = G.argmax(axis=1) # BFM 近似解 j_bfm, _, _ = bfm_map(X, Y, iters=300) # 对比传输代价 cost_exact = (G * M).sum() / N cost_bfm = M[np.arange(N), j_bfm].mean() print(f"精确代价 {cost_exact:.6f} / BFM 代价 {cost_bfm:.6f}")

在我这边跑出来两个代价差在 1% 以内,映射一致率大约 85%~90%。注意这是“近似一致”而不是完全相同,因为离散点云上最优映射可能不唯一(尤其是点稀疏的地方),BFM 收敛到的解和线性规划给出的解在退化情形下可以不同,但传输代价应该接近。这个对比的意义在于:如果 BFM 的代价明显高于精确解,说明迭代没收敛或者实现有 bug;如果只差一点,就可以放心把它塞进 JKO 循环里了。

3.4 参数选择:τ、N、迭代次数与收敛判据

参数这块我踩坑最多,单独拎出来讲。τ 是 JKO 的时间步,它的作用和显式格式里的步长完全相反——τ 越大,单步的隐式性越强,稳定性越好,但离散误差越大,因为 JKO 的解和连续流的偏差是 O(τ) 量级。τ 越小越接近真实流,但需要的步数变多,而且每步的传输问题在 τ→0 时趋向恒等映射,数值上反而不太有信息量。我的经验是先用 τ=0.02~0.05 跑通,观察能量曲线是否光滑,再根据精度需求调整。

粒子数 N 决定了密度表示的精细程度。一维用几千个粒子足够了,二维想要看清楚的密度结构,一般要上万。但 N 上去之后距离矩阵是 O(N²) 的,N=10000 时单精度也要 400MB,这时候必须换 k-d 树或者分块。我的做法是:N ≤ 3000 直接用矩阵,超过就用scipy.spatial.cKDTree做近邻查询,把 c-transform 改成“查询若干个近邻再取最小”,牺牲一点精确性换内存。

BFM 的迭代次数和收敛判据也要定。我的判据是前向映射索引在连续 5 轮内变化比例低于 0.1%,或者达到最大迭代次数。加这个判据之后,平均迭代次数从固定的 300 降到 60 左右,整体速度快了三四倍。

参数典型取值调大的后果调小的后果
τ(JKO 步长)0.02~0.05离散误差 O(τ) 明显,粒子位置偏移需要步数暴增,传输问题趋近恒等
N(粒子数)1e3~2e4内存 O(N²) 爆炸,单步变慢密度分辨率不足,映射不唯一
BFM 迭代次数50~300收益递减,浪费时间映射未收敛,传输代价偏高
阻尼 θ0.3~0.7接近 1 时震荡过小则收敛缓慢

还有一条是正则化。如果你要做的是熵正则最优传输(Sinkhorn 那套),会额外冒出正则参数 ε,它的取值和 τ、N 都有关联,ε 太小数值不稳定,太大则偏离真实传输。这一块和 BFM 的原始形式是两条路线,我这里不展开,避免混淆。

4. 踩坑记录:Windows 上最常见的问题与排查路径

下面这些坑基本都是我在 Windows 上真实撞过的,按“安装期”和“运行期”分开列,附上我的判断思路。表格先给速查版,后面挑几个典型展开。

4.1 安装期报错速查表

报错关键字大概率原因处理方式
无法加载文件 ...profile.ps1,因为在此系统上禁止运行脚本PowerShell 执行策略限制Set-ExecutionPolicy -Scope CurrentUser RemoteSigned
ModuleNotFoundError: No module named 'ot'包名写错,装成了别的包重装pip install POT,import 名才是ot
OMP: Error #15: Initializing libiomp5md.dll多个 OpenMP 运行时共存统一用 conda 装包,或临时设KMP_DUPLICATE_LIB_OK=TRUE
Microsoft Visual C++ 14.0 or greater is required缺 C++ 生成工具装 Build Tools 的 C++ 桌面开发组件
DLL load failed while importing二进制不兼容 / 位数不一致重建环境,确认全是 64 位
OSError: [WinError 126] 找不到指定的模块依赖的 DLL 不在 PATH用conda install补依赖,别手工拷 DLL
UnicodeDecodeError: 'gbk' codec can't decode默认编码是 cp936读文件时显式encoding="utf-8"
pip下载超时网络原因换国内镜像源,或延长 timeout
FileNotFoundError路径里有奇怪字符中文/空格路径项目放纯英文路径,别放 OneDrive

OMP: Error #15这个值得单独说。它的本质是程序加载了两份 OpenMP 运行时,各自以为自己是唯一的,初始化时互相打架。临时解法是设环境变量KMP_DUPLICATE_LIB_OK=TRUE让它闭嘴继续跑,但这个做法有风险——两套运行时同时管理线程池,可能出现结果不稳定甚至崩溃,而且性能会掉。我在调试阶段会用它应急,正式跑实验前一定把冲突的包统一掉。判断方法是在环境里跑conda list | findstr mkl和conda list | findstr openmp,看看是不是混着装了两套。

还有个隐形的坑:Set-ExecutionPolicy改完之后需要重开终端才生效,我第一次改完直接在原窗口试,还是报同样的错,白折腾十分钟。

4.2 运行期数值异常:能量不降、粒子塌缩、步长发散

安装问题好查,数值问题才磨人,因为不报错,只是结果不对。我遇到过的典型症状有这么几类。

能量不单调下降。理论上 JKO 保证能量单调不增,如果出现上升,先查 τ。τ 太大时,单步的隐式格式虽然稳定,但传输映射的近似误差会累积;把 τ 减半再跑,如果曲线变光滑了,那就是步长问题。如果减了还不行,检查梯度函数的符号——我遇到过把gradV写成-gradV的情况,能量当然一路往上飙。再不行就看求解器是否收敛,L-BFGS-B 返回的res.success是不是 False,很多人不检查这个标志,结果拿着未收敛的解当结果。

粒子塌缩到一点。这个症状在只有外势能的时候正常,因为二次势能确实会把所有粒子拉到原点。但如果你加了熵项或者相互作用项还塌缩,那就是能量项没生效。检查方法很简单:单独跑一步,打印能量各分量的数值,看熵项的贡献是不是零。

距离矩阵内存爆掉。N=20000 时距离矩阵是 4×10⁸ 个元素,float64 就是 3.2GB,直接 MemoryError。改成cKDTree查询,或者分块算——按行切 500 一块,每块算完取 argmin 再释放。分块版本我实测比一次性算慢 20% 左右,但内存占用从 GB 级降到百 MB 级。

多进程无限递归。这个问题在 Windows 上极常见。Linux 的 fork 会把父进程内存复制一份,子进程直接能用;Windows 的 spawn 会启动一个新的 Python 解释器,重新导入主模块。如果你的并行代码写在模块顶层没有if __name__ == "__main__":保护,子进程导入模块时又会创建子进程,一路递归下去,最后报RuntimeError: An attempt has been made to start a new process before the current process has finished its bootstrapping phase。解决办法就是老老实实加保护,把并行入口包起来。这个坑我在参数扫描脚本上踩过一次,那次任务管理器里刷出几百个 python 进程,机器直接卡死。

# 错误写法:并行代码裸露在模块顶层 from concurrent.futures import ProcessPoolExecutor with ProcessPoolExecutor() as ex: results = list(ex.map(run_one, params)) # 正确写法 def main(): with ProcessPoolExecutor() as ex: results = list(ex.map(run_one, params)) if __name__ == "__main__": main()

4.3 性能调优:向量化、并行与内存

跑通之后自然会想快一点,这里分享几个实测有效的做法。

把 Python 循环换成向量化。c-transform 那部分,如果写成for i in range(N): for j in range(M):的双重循环,N=M=1000 就要跑一百万次,纯 Python 大概几十秒。改成 numpy 的广播运算,同一规模不到 0.1 秒,差了两个数量级。代价是内存,所以规模上去后要配合分块。

用 cKDTree 代替全量距离矩阵。半离散最优传输里,每个源点只需要找最近的若干个目标点,全量算矩阵是浪费。cKDTree的query支持workers=-1用满所有核心,在 Windows 上这个参数是可用的,实测 8 核机器上有 4~5 倍加速。

from scipy.spatial import cKDTree import numpy as np tree = cKDTree(Y) # 查 8 个近邻,workers=-1 用满核心 dist, idx = tree.query(X, k=8, workers=-1)

并行策略选对粒度。参数扫描适合进程级并行,因为每个任务独立;单次迭代内部的循环适合线程级并行,但 numpy 本身已经用了多线程 BLAS,你再套一层线程池往往没有收益,甚至会因为争抢 GIL 和线程调度变慢。我的经验是:外层任务用ProcessPoolExecutor,内层计算交给 numpy,不要在中间再插一层线程池。

注意首次导入的冷启动。前面提过 Windows Defender 的影响,另一个因素是 numpy 首次 import 要加载几十个 DLL。我在做参数扫描时,每个进程都要重新 import 一次 numpy,如果任务本身很短(比如几十毫秒),冷启动开销可能比计算还大。这种情况下要么把任务打包成大块,要么用initializer预热。

5. 几句实操心得,给准备入坑的人

我在 Windows 上折腾这套东西最大的感受是:数学上的难点和工程上的难点完全是两回事。BFM 的 c-transform 迭代看一遍就懂了,真正花时间的是环境变量、编码、多进程保护、内存这几件事。所以如果你的目标只是快速验证一个想法,别一上来就追求把论文里的完整格式实现出来,先用一维的闭式解把环境和基本逻辑锁死,再逐步上二维、上复杂能量项,这个渐进路径能帮你把问题范围缩到最小。

另一个体会是,不要跳过对照验证。我现在写任何传输相关的代码,都会先拿 POT 的精确解对一遍代价,差得离谱就说明有 bug,差一点就说明只是近似。这个习惯帮我省了无数个小时。同样地,能量单调性、粒子保

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进