作者 | CW不要無聊的風格 编辑 | 极市平台
原文链接：https://zhuanlan.zhihu.com/p/696121622/
点击下方卡片，关注“自动驾驶之心”公众号
戳我-> 领取自动驾驶近15个方向学习路线
本文只做学术分享，如有侵权，联系删文
导读
来看看S-DSB到底是个什么东西！
邂逅
前段时间的某个早上，在“晨例”喝手冲(咖啡)的时候不小心看了场介绍 “简化薛定谔扩散桥(Simplified-Diffusion Schrödinger Bridge)” 的直播，是 paper(https//ar5iv.labs.arxiv.org/html/2403.14623) 原作的技术分享，感觉挺有意思的，而且近来这种基于扩散薛定谔桥(DSB)的框架貌似蠢蠢欲动想接替扩散模型(Score-based Generative Model)成为视觉生成界的新宠。于是在这种契机之下，CW 就没忍住把 S-DSB 的 paper 和 源码实现都看完了..
你们也知道 CW 的个性，看完了不无聊的 paper 通常都忍不住要吹水一番；而对于开源的 paper 也忍不住要亲手撸一遍，总忍不住把好玩的东西都内化到自己身上留下烙印。唉~坏毛病——又得费精力码字了。
简化扩散薛定谔桥(S-DSB)是针对扩散薛定谔桥(DSB)(https//arxiv.org/abs/2303.16852)收敛慢和训练难的问题进行改进而提出的，主要的改进措施包括：
简化了 DSB 的 loss 函数
将 SGM 适配到 DSB 框架
提出了两种重参数化方式
话又说来，先不论 S-DSB 是怎么回事，那所谓的什么 DSB 和薛定谔的猫到底是什么关系呀？它和 SGM 相比又是怎样的一种玩法呢？
嗯，好问题！DSB(扩散薛定谔桥) 中的“薛定谔”，确实是薛定谔的猫的主人，大名鼎鼎的量(浪)子物理学家——薛定谔(https//baike.baidu.com/item/%25e5%259f%2583%25e5%25b0%2594%25e6%25b8%25a9%25c2%25b7%25e8%2596%259b%25e5%25ae%259a%25e8%25b0%2594/2124805)。大家也知道，他是搞量子力学的，这么一来 DSB 就像是个“跨界er”—— 由量子力学跨界来搞 CV 生成。这..是要来一波降维打击和砸场子？
同时，在当下，基于 score 的扩散生成模型 —— SGM 可谓中流砥柱，但这在一定程度上也得益于大家都很给面子。
毕竟嘛.. 它的毛病显而易见：
先验被限制为简单的有解析形式的分布(如高斯分布)而不能是任意分布
只能单向生成(由先验向目标数据分布)，而不能在任意两个分布之间建立双向映射
然而，DSB 在先天上就不存在这些毛病，这很降维打击！并且它的目标是计算任意两个可采样分布之间的最优传输(Optimal Transport,OT)(https//www.damtp.cam.ac.uk/research/cia/files/teaching/Optimal_Transport_Notes.pdf)，有 OT 理论的加持。注意是“任意”分布！并没有将分布限制为某种形式或要求具有解析式。一旦完成训练，它就可以完成双向生成——将 A 分布的数据喂给模型，就能采样生成 B 分布的数据，反之亦然，并且 A, B 可以是任意分布，这天花板可高多了！
由此也可以发现，DSB 天生就是条件生成的优秀选手，比如在训练时给的两个分布是猫和狗的图片，那么训练完毕后它自然就可以根据一张猫的图片，生成一张“与这只猫很像的狗”的图片(反之亦然)，不需要额外为“根据猫来生成狗”这个条件而进行设计。或者说，在 DSB 这就不存在 conditional generation，因为它天生就是 conditional 基因。
WOW!! 这简直就是不无聊的风格，很符合 CW 的口味！
不多说了，接下来就切换到稍微正经的模式来进入到本文的世界中吧~
【注】
本文会反复使用一些简称，主要有：
SB —— 薛定谔桥
DSB —— 扩散薛定谔桥
S-DSB —— 简化扩散薛定谔桥
SGM —— 泛指扩散模型
S-DSB 诞生的源动力
一个新事物的诞生通常源于当下时代背景的一些因素，但因素众多，不可能逐一分析，于是作者选择抓“主要矛盾”——精准定位到 SGM 和 DSB 这俩家伙身上，然后仔细挑遍了它俩的毛病，经过筛选后，最终列出的“重大症状”如下：
SGM
针对不同任务需要很费心思去设计加噪方案(noise schedule)
先验(prior)被限制为一种预定义的简单分布(e.g. Gaussian)
DSB
模型收敛慢、训练难、拟合能力不足
对空间和计算资源的要求高，在训练过程中，网络学习的 target 需要经历两次独立的
forward()才能获得
另外，DSB 没有利用好 SGM 当前的发展成果，相当于有很好的资源却浪费了。
是病就得治，为了让打造健康和谐的生成模型生态，作者提出了一种 DSB 的改进版本，它就是本文的主角—— S-DSB(Simplified-DSB)，这个新玩法对原来 DSB 的 loss 函数进行了简化，并且将 SGM 拿到 DSB 框架中去玩(bridge the gap between DSB & SGM)，同时还可以利用预训练的 SGM 做初始化，从而加快收敛。这种方法不仅解决了以上 DSB 的问题，还进一步提升了原有 SGM 的性能。
此外，作者还提出了两种重参数化技巧(reparameterization trick)应用到 DSB 的输出空间，从而对应到不同的训练目标，进一步提升了模型的收敛速度，并且肉眼可见比 DSB 生成的效果要好。
这还不止，作者设计实验进一步证明，即使在没有使用预训练 SGM 的情况下，S-DSB 也能取得与 DSB 不相上下的性能。
其实作者大大你是不是想说：我 S-DSB 完爆你 DSB！
【配图】
前情回顾
在主角 S-DSB 正式登场之前，先来回顾下历史情节—— SGM 和 SB&DSB 的玩法，以更好地承接剧本。
SGM 的玩法
SGM 通过两个互逆的过程(dual process)将两个分布 “连接” 了起来：前向过程(forward process) 将数据分布 转换为先验分布 ，而逆向过程则相反。这两个过程都建模为马尔科夫链，在离散时间步的情况下，联合分布可表示为:
但是，直接根据上面的贝叶斯定理去计算逆向过程的 通常是 intractable 的，毕竟 都不知道。因此 将前向过程建模为不断施加高斯噪声的马尔科夫链，使得:
其中， 起到了 noise scale 的作用。然后，再通过贝叶斯定理，就能得到逆向过程的 也近似等同于高斯分布。不 show 明白你们肯定不服，好吧，具体推导如下：
然后对 在 处进行1阶泰勒展开:
将上式代入到前面的式子中，得到:
同理，这种玩法可以拓展至连续时间的情况，从而前向过程就对应到以下 SDE：
其中 代表布朗运动， 是漂移(drift)项， 是扩散(diffusion)项。这样，如 (i) 式的 transition kernel 实质上就是对上式使用欧拉-丸山(Euler-Maruyama)法进行数值离散的结果。
类似地，逆向过程则对应到以下 reverse-time SDE：
注意，上式中 也是正的，而在通常所见的 reverse-time SDE 中， 是一个负值的微元。
SB 和 DSB 的玩法
既然 DSB 与 SGM “同病相怜”，那么也同样先来回顾下它的玩法。然而考虑到 DSB 是在 SB 基础上发展而来的，所以干脆就先从 SB 开始下手叭~
SB
薛定谔桥(Schrödinger Bridge) 本是数学和物理学中的一个概念，起源于量子力学中的薛定谔方程，同时也多应用在最优运输问题中。它本质上是一种条件随机过程，通常指在给定起点和终点分布的条件下，对某种随机过程(e.g. 布朗运动)的变换。
恰好 SGM 玩的就是典型的起点和终点分布问题——前向过程是复杂的数据分布到简单的高斯分布，逆向过程则相反。于是，把 SB 引入 SGM 的游戏中就成了很顺手的事情。
薛定谔桥问题是在一定约束条件(i.e. 给定起点和终点)下的最优传输问题，其寻求在两点所有可能的轨迹中，使得 cost 最小的那一条。若将上一章介绍的 作为目标轨迹，将其记为 “参考分布" ， 则 的目标就是要找到一个分布 ， 满足:
待找到这样的 之后，就可以使用祖先采样(ancestral sampling)由 开始，不断根据 来采样得到 ；同理，也可以反过来由 开始，不断根据 来采样得出 。
SB 是求解一个可逆的双向映射，相比起 SGM，它没有将其中一个分布预设为已知解析形式的简单分布(SGM 的做法通常是设置 )。也就是说，SB 的目标是实现任意两个分布之间的转换。
在大多数情况下，SB 的目标都不好求(且通常没有解析形式)，于是为了降低求解难度，经常会用到一种叫作 "Iterative Proportional Fitting(IPF)"(https//www.sciencedirect.com/science/article/abs/pii/016771529390257J%3Fvia%253Dihub) 的招数，它本质上属于一种弱化了边界条件的迭代优化方法：
其中 代表迭代的轮数、 。可以看出，在每一轮优化中，仅有 或 这其中一个边界条件，而不像原始 做法那样需要同时满足两者。以这种方式，当每一轮都取得最优解并且迭代轮数足够大时，最终将得到 的目标解 。
然而，在每轮优化中，这种方法仍需要计算和优化联合概率分布 ，在面对高维数据时对计算和内存资源就会有很高的要求。
DSB
联想到扩散模型 经常 “玩弄” 和 ，也就是在相邻时间步之间的条件分布，于是 DSB 就汲取了这个灵感来近似 IPF 的做法， 它将原本在每轮中优化联合概率分布拆解为优化一系列相邻时间步之间的条件分布，从而将 在前向过程(forward process)中解构为一系列 ；在逆向过程(backward process)中则解构为一系列 ，然后在训练过程中交替地去优化 forward 和 backward 两条轨迹:
如上，在奇数轮(epoch) 优化的是 ，也就是 所对应的 backward 轨迹，得到 ；而在偶数轮 优化的则是 ，对应 的 forward 轨迹，得到 。每轮在优化时都最小化与上一轮所得结果的 KL 散度。
既然都向 SGM 汲取了灵感，不如将它的建模方法也借(白)鉴(嫖)过来好了。
于是，DSB 就将 和 都建模为高斯分布，这样两者都拥有解析形式，方便求解和计算。同时，分别使用两个模型去建模 forward 和 backward 轨迹。经过一系列花哨的数学推导 (细节请参考 DSB 的 paper)，最终得到 loss 函数如下:
在实现时，如果当前处于第 个 epoch，且 为奇数，则优化的是 backward 轨迹，建模 backward 过程的模型就会被训练来去估计 ，也就是 backward 高斯分布 的均值；类似地，若 为偶数，则优化的是 forward 轨迹，建模 forward 过程的模型会被训练为去估计 ，对应 forward 高斯分布 的均值。
师出同门：各生成模型实质是扩散薛定谔桥的特例
在 S-DSB 正式亮相之前还有个彩蛋——作者揭出目前市面上流行的一些生成模型如 SGM, FM(Flow Matching)(https//arxiv.org/abs/2210.02747) 等都可看作扩散薛定谔桥的特例。这是因为，之前已经有大佬的 paper(https//arxiv.org/abs/1308.0215) 指出薛定谔桥的最优解实际上遵循以下形式的 SDE：
更为重要的是， 在 时刻的分布 就是:
通过与前面所介绍的 SGM 的 SDE 形式相对比，容易看出 SGM 实际上就是式(11)在 是常数(从而 ) 且 时的特例。根据 在 时刻的分布(即上式)，这同时有:
将其代入到上图式(11)中的第二个式子则正好就是我们所熟悉的 SGM 的 reverse-time SDE(这里的 reverse-time SDE 中 为负，而前面章节所展示的是 为正的形式)!
进一步来说，只要设置不同形式的 和 ，就能够对应到不同生成模型的玩法。比如在 SDE 统一扩散模型的那篇 paper 中所提出的 Variance Preserving (VP) 和 Variance Exploding(VE) 的加噪方案(noise schedule)，就是将漂移项 设成 ，其中 是非负实数。
主角登场：S-DSB
主角 S-DSB 登场了！请大家倾耳拭目，一起来看看 S-DSB 是怎么为它前辈 DSB 做手术治病的。
手术一：简化 loss 函数
S-DSB 的命名由来主要是源于它对 DSB 的 loss 函数进行了简化，因此拿了 "simplified" 的头衔。具体地，它将 DSB 的 loss 函数简化为以下形式：
在 和 的假设条件下，S-DSB 的这个简化版 loss 函数可近似等于原来 DSB 的 loss 函数，比如对于 backward model 的 target，容易证明当 时 ， 即两个 loss 函数的 target 近似相等:
那么，这个简化版的 loss 有何优势或者说能够带来什么好处呢？当然是有的，不然作者也不敢写这篇 paper 来讲故事。
概括起来，S-DSB 的简化版 loss 主要有以下两个好处：
直观来看，这种 loss 形式就是要网络去学会预测 “下一个状态” —对于 backward model 来说，它要学会基于 去预测 ，其轨迹方向是 ，在当前是 的情况下，下一个状态就是 ；而对于 forward model 来说，它则要学会基于 去预测 。这种 loss 形式既直观又易于理解，并且与 SGM 实现了 “对齐” (SGM 所用的 loss 通常也是这种形式)。
更为重要的是，相比于原来 DSB 的 loss，这种简化版 loss 在计算网络学习的 target 时不必运行两次模型的
forward()过程，从而节省了计算和内存资源。比如对于 backward model 来说, DSB 的 loss 需要运行 ，这是2次模型的forward()，而在这里网络的 target 就直接是 。另外，模型运行forward()的次数也被成为 NFE-Number of Forward Evaluations.
手术二：让预训练 SGM 助攻一波
虽然 S-DSB 对 loss 函数进行了简化，但它的训练方式还是处于 IPF 框架之下的。也就是说，它需要交替训练两个 model，在训练其中一个时，其学习的 target 轨迹实际上由另一个 model(固定住权重) 给出，因此每一轮的训练效果实质上会依赖于上一轮的收敛性！作者在 paper 的 3.2 节也给出了一小段数学解释，感兴趣可以去瞄瞄，挺直观易懂的~
另外，通过对原有 DSB 的 loss 函数进行简化，S-DSB 的 loss 已经与 SGM 实现了对齐，这就代表，S-DSB 的训练目标与 SGM 其实是一致的。
炼过丹的都知道，当遇到训练比较难收敛时，如果有预训练的权重，那么想都不用想当然是利(白)用(嫖)起来啊！
目前我们的 S-DSB 正是面临这种处境，刚好市面上又大把已经训好的 SGM，而 S-DSB 的训练 target 又与 SGM 一致，这..八字也太合了叭，你要不加载预训练的 SGM 权重就简直是违反天时地利人和！
正是基于这种先天条件与后天境遇的结合，作者决定让 pre-trained SGM 来助攻一波，将参考分布 根据 pre-trained SGM 的 noise schedule 进行设置就可以使用它来初始化 backward model; 同理，forward model 也是可以用 pre-trained SGM 来进行初始化的。两个模型可以选择不同的预训练权重，也可以选择一样的，至于效果嘛..还得看具体场景，毕竟是炼丹=玄学
手术三：重参数化(Reparameterization)
除了简化 loss + 白嫖预训练 SGM 以外，S-DSB 还有招，那就是 SGM 的网红招数——重参数化。又白嫖..？
这招通常能够降低模型的训练难度，因为模型学习的 target 通常是随时间动态变化的，这又源于整个过程本质上服从 SDE。于是，让模型直接去预测这种动态变化的 target，通常比较困难，从而需要更长时间才能收敛。经过重参数化后，通常能够使模型在不同 timestep 都拥有一致的 target。
作者大大当然也晓得这种网红技巧带来的好处，而且如果能用上这么 in 的 trick，多少也能让 paper “漂亮”几分。
于是，其果断决定在 S-DSB 中引入重参数化技巧，而且还给重参数化后的 S-DSB 起了个新名字 Neparameterized DSB -DSB)”。OK，决心是下了，但实际该如何行动呢?
要使用哪种重参数化方式，得考虑决定因素有哪些。如果在 backward 或 forward 轨迹中，始终有一致性的状态(比如它们的终点 或 ) 决定着整条轨迹的变化，那么可以考虑让模型学会预测终点就好了，这样在每个 timestep，模型的 target 都是一致的。
顺着这种 “幻想"，作者还真捣鼓出了不得了的东西:
(证明过程请参考论文附录 B.1~B.2，写得非常详细，而且没有过份跳跃，顺着看下来能够明明白白，赞一个！)
以上提出的这个命题说明：在学习 forward 轨迹时，模型需要预测的下一状态 其实是由当前状态 和 forward 轨迹的终点 所决定的; 类似地，在学习 backward 轨迹时，模型需要预测的 是由前一个状态 和 backward 轨迹的终点 所决定的。
嗯哼? 这说明什么，是不是有点 feel 了
Terminal Reparameterized DSB(TR-DSB)
不卖关子了(相信聪明的你们也想到了)，上面那个命题其实已经指明了道路——让模型学会去预测终点，因为无论在哪条轨迹上(forward or backward)，模型所要预测的下一个状态都是由当前状态和终点所共同决定，而当前状态是已经拿到手的，所以只要模型能够预测出终点值，那么就可以由当前状态和模型所估计的终点值来计算出下一个状态。不断迭代这个过程最终就能抵达轨迹的终点，从而完成整个采样过程，这本质上属于祖先采样的方式。
OK，既然重参数化后，模型预测的目标政变了，那么就得修改原来 S-DSB 的 loss 函数——将 backward model 的 target 改为 ；同时将 forward model 的 target 改为 :
这种重参数化的方式是基于终点(terminal)的重参数化，于是在这种模式下的 S-DSB 就被称作 "Terminal Reparameterized DSB(TR-DSB)"。
待训练完毕后，对于 backward model，我们就可以将模型的输出看作是 ，代入到命题 3 的 (19)式中的第一式从而完成 的采样过程; 同理，对于 forward model，则将其输出看作是 代入到命题 3 的(19)式中的第二式来完成 的采样生成。
Flow Reparameterized DSB(FR-DSB)
然而，基于命题 3 的重参数化方式也并非仅有预测终点这一种，仔细观察(19)式，不难发现其实还可以令 backward model 去预测 ; 同时 forward model 则去预测 ，对应到以下 loss 函数:
这种预测目标并非是强行瞎搞的，而是有直观意义的：对于 backward 轨迹 来说， 代表由当前状态指向终点的向量，而 则代表在 forward 的加噪过程中， 的噪声尺度总和，相当于变化了多少，我们可以把它当作 与 的 “距离” ；同理，对于 forward 轨迹来说， 代表由当前状态指向终点的向量， 则同样看作 与 之间的 “距离”。
于是，这种预测目标所代表的意义就是由当前状态指向终点的“单位”向量。
这种单位向量可看作“流(flow)”，并且，这种预测目标是受到了 Flow Matching 的启发(白嫖 again)，于是这种玩法下的 S-DSB 也被称作 "Flow Reparameterized DSB(FR-DSB)"。
当模型训练完毕要用于采样生成时，比如 backward model，则直接将命题3的(19)式中的 替换为模型的输出 ，就可以得到下一个状态 (当前状态是 ) 的采样均值。随后不断迭代这个过程并遵循祖先采样的套路，就可以完成整个采样过程。对于 forward model 也是同理， CW 就不再啰嗦了
ps: 对于以上两种重参数化方式，作者在论文附录 B.3 部分有更进一步的数学解释，讲得挺详细的，感兴趣的友友们可以去 look look~
Noise Scale
Noise scale 代表着所施加的噪声强度。由于在 DSB 的框架下，所谓的 可能也是有意义的目标数据分布因此无论是 还是 的过程，噪声强度都不能一昧地单调变化，毕竟两端都是有意义的目标分布，于是作者就采取了一种 “对称” 方案，即先增大、后减小，增减幅度一致：
另外，作者进一步发现，使用尺度(数量级)更小的 noise schedule 在由源分布转换至对端分布的过程中会倾向于保留更多源分布的信息(比如姿势、色彩等)；而尺度更大的 noise schedule 则可能产生更高质量(high-quality)的生成效果(即生成的图片与对端分布中的更像、看起来更自然与逼真)。
生成模型的通用训练框架
作者发现，包括 DSB 在内的大部分生成模型实际共享着同一套训练框架(逻辑)，只是在训练目标和加噪方式上存在差异。为此，作者将它们的训练逻辑总结为如下：
其中每个时间步的状态(隐含着加噪方式)以及对应的 target 如下：
进一步，作者给出了 S-DSB 的训练 pipeline：
其中，第1个 while 循环代表 backward model 的训练逻辑，它要基于时间步 以及对应的状态 来预测状态 。需要注意的是，它的输入状态 是由 forward model(假设已经收敛) 的估计值 加上噪声项 所 “制造” 出来的， 这是由命题3以及 DSB 对 forward 过程的建模定义(即高斯分布)所决定的。
第2个 while 循环则代表 forward model 的训练逻辑，它要基于时间步 及对应状态 来预测状态 。同理，其输入状态 是由 backward model(假设已经收敛) 的估计值 加上噪声项所计算出来的，这同理也是由命题3以及 DSB 对 backward 过程的建模定义所决定。
进一步展开说说，之所以 training model 的输入状态要由另一个 model 给出，是因为命题3指出 training model 所学轨迹中的条件分布(e.g. ) 实际上是另一个 model 所建模轨迹中的后验分布(e.g. )。
同时，DSB 对于 forward 和 backward 过程都建模为高斯分布，其对应马尔科夫链的 transition kernel 表现为以下形式：
forward process
backward process
在训练 backward model 时，就假设 forward model 已经收敛，于是将后者的输出 当作 ，也就是 的均值，在训练 forward model 时也是同理。
之所以对这部分这么啰嗦，其实是为了在接下来的代码实现环节中能够有更清唽的逻辑，真的不是 CW 想要强行吹水。
玩一把：从 0~1 实现 S-DSB
经过前面的铺垫，现在终于可以动手玩玩 S-DSB 了！在本章，CW 会和大家一起从 0~1 实现 S-DSB，主要包括分成以下几部分：
noise schedule
DSB 建模
神经网络
数据集
训练 pipeline
Noise Schedule 的实现
首先来实现 noise schedule，它决定了两个分布之间的数据是如何进行转换的，等同于 SGM 的加噪方法。
沿用 SGM 的惯例，我们将两个分布分别记为 ，并且规定 时数据服从 分布； 时数据服从 分布，记数据为 ，那么就有 。在时间 的过程中 ，也就是本属于 中的数据会逐渐演变为服从 的数据，而 noise schedule 就在数学上决定了 是如何一步步变成 的。
既然只是简单玩玩，不妨选择最简单直接的方案一一线性插值。记数据在时间 时为 ，那么对于每对起点和终点: ，就有 ，这样正好有:
在这个线性插值的方案中， 的系数分别是 ，于是在下面的代码实现中 就将 命名为 coeff 0。；将 命名为 coeff_1，这两个变量也对应于论文中的 。另外，在代码实现中， 会将 表示为 x_1 变量。
对于 noise scale ，我们就遵循 paper 中的方案，让它先线性递增、后线性递减，增减速率相等，呈对称性，用 prepare_gammas() 函数来封装这个逻辑。
以上逻辑都包含在 FlowNoiser class 中， 如下：
import torch
import torch.nn as nn
from typing import Dict, Tuple, Union
from torch import Tensor
from torch.nn import Module
def align_shape(x: Tensor, coeff):
""" 将 coeff 的维度与 x 对齐. """
if isinstance(coeff, dict):
for k, v in coeff.items():
if isinstance(v, Tensor):
while len(v.shape) < len(x.shape):
v = v.unsqueeze(-1)
coeff[k] = v
elif isinstance(coeff, Tensor):
while len(coeff.shape) < len(x.shape):
coeff = coeff.unsqueeze(-1)
return coeff
class FlowNoiser(Module):
def __init__(
self, device,
training_timesteps: int, inference_timesteps: int,
gamma_min: float, gamma_max: float,
simplify: bool = True, reparam: str = None
):
""" 采取流式的 noise schedule, 实质就是在两个目标分布之间线性插值. """
super().__init__()
self.device = device
# noise scale 的取值范围
self.gamma_min = gamma_min
self.gamma_max = gamma_max
# 训练阶段的时间步
self.training_timesteps = training_timesteps
self.training_timestep_map = torch.arange(
0, training_timesteps, 1,
dtype=torch.long, device=device
)
# 推理阶段的时间步
self.inference_timesteps = inference_timesteps
# 当推理时的时间步数与训练时不一致时,
# 可以根据步长映射至训练时的时间步
self.inference_timestep_map = torch.arange(
0, training_timesteps, training_timesteps // inference_timesteps,
dtype=torch.long, device=device
)
self.num_timesteps = training_timesteps
self.timestep_map = self.training_timestep_map
# 是否是 S-DSB(否则就是 DSB)
self.simplify = simplify
if simplify and reparam is not None:
if reparam not in ("FR", "TR"):
raise ValueError(f"reparam must be 'FR' or 'TR', got: {reparam}")
# 所使用的重参数化方式: flow or terminal
# 在 S-DSB 的情况下才生效
self.reparam = reparam
def train(self, mode=True):
self.num_timesteps = self.training_timesteps if mode else self.inference_timesteps
self.timestep_map = self.training_timestep_map if mode else self.inference_timestep_map
def eval(self):
self.train(mode=False)
def coefficient(self, t: Union[Tensor, int]) -> Dict:
""" 用于在两个目标分布之间进行插值的系数. """
if isinstance(t, Tensor):
t = t.max()
if t >= len(self.timestep_map):
coeff_0, coeff_1 = 0., 1.
else:
timestep = self.timestep_map[t].float()
coeff_1 = timestep / self.training_timesteps
coeff_0 = 1. - coeff_1
return {"coeff_0": coeff_0, "coeff_1": coeff_1}
def prepare_gammas(self) -> Tensor:
""" 使用线性对称的 noise scale,
也就是先线性递增至最大值; 后线性递减至最小值. """
gammas = torch.linspace(
self.gamma_min, self.gamma_max,
self.num_timesteps // 2,
device=self.device
)
self.gammas = torch.cat([gammas, gammas.flip((0,))])
def forward(self, x_0: Tensor, x_1: Tensor, t: Union[Tensor, int]) -> Tensor:
""" 在两个目标分布之间进行插值, 从而得到中间各时间步的状态. """
coeff = align_shape(x_0, self.coefficient(t))
coeff_0, coeff_1 = coeff.values()
x_t = coeff_0 * x_0 + coeff_1 * x_1
return x_t
@torch.no_grad()
def trajectory(self, x_0: Tensor, x_1: Tensor) -> Tensor:
""" 生成由 x_0 至 x_1 的轨迹. """
trajectory = [x_0.clone()]
for t in range(self.num_timesteps):
x_t = self.forward(x_0, x_1, t)
trajectory.append(x_t.clone())
return torch.stack(trajectory)FlowNoiser 被实现为 torch.nn.Module 子类，这是为了方便其在训练和推理之间切换所使用的 timestep，因为训练和推理所使用的总时间步数可能是不一样的，而 torch.nn.Module 有 train() 和 eval() 模式，这样就可以借助这两个模式的切换来对应到训练和推理过程所用 timestep 的切换。
至于 forward() 方法，就是根据给定的 x_0 和 x_1 以及当前时间 t 按照线性插值方案计算出 x_t；trajectory() 方法则负责记录整条 forward 轨迹，即收集从 x_0 演变至 x_1 过程中的所有状态。
由于第1个 epoch 训练 backward model 时，forward model 还未训练，因此前者的输入状态需要由 FlowNoiser 给出。参考 Algorithm 2，为了和 forward model 收敛时所给出的输出对齐，那么就要让 FlowNoiser 能够给出 ，它是 的均值: ，于是我们就实现一个叫作 forward_F() 方法来计算这个值。
输入状态有了，ground truth(仪练目标) 也要对应给出，它有四种情况：原始 DSB、S-DSB(不使用重参数化)、TR-DSB 以及 FR-DSB，每种情况根据论文公式去计算即可。哦，别忘了时间步也要输入到模型中。
这么一来也就是说，在每个时间步，我们都需要 (输入状态、时间步、训练目标) 这三元组，CW 于是实现一个叫 trajectory_F() 的方法来收集 过程中的所有三元组。
def forward_F(self, x: Tensor, x_1: Tensor, t: Union[Tensor, int]) -> Tensor:
""" DSB 的 F_k^n 函数, mean(x_{k+1}) = F_k^n(x_k)
在 S-DSB 中被建模为 q 的后验分布, 见论文公式(19)第二式. """
coeff_0_t = align_shape(x, self.coefficient(t))["coeff_0"]
coeff_0_t_plus_one = align_shape(x, self.coefficient(t + 1))["coeff_0"]
vec = (x_1 - x) / coeff_0_t
F_x = x + (coeff_0_t - coeff_0_t_plus_one) * vec
return F_x
@torch.no_grad()
def trajectory_F(
self, x_0: Tensor, x_1: Tensor,
sample: bool = False
) -> Tuple[Tensor]:
""" 根据 DSB 建模的 forward process(F_k^n 函数) 计算出
x_0 -> x_1 的整条轨迹. 这也同时可作为 backward model 在
首个 epoch 里的训练 target. """
self.prepare_gammas()
ones = torch.ones((x_0.size(0),), dtype=torch.long, device=self.device)
x = x_0
x_all, gt_all, t_all = [], [], []
for t in range(0, self.num_timesteps):
''' 收集轨迹的各时间步 '''
t_ts = ones * t
t_all.append(t_ts)
''' 收集轨迹的各状态 '''
# mean(x_{k+1}) = F_k^n(x_k)
F_x = self.forward_F(x, x_1, t_ts)
# 如果是采样过程的最后一步, 则不加上噪声项
# 理论依据是 Tweedie's formula
if sample and t == self.num_timesteps - 1:
x_next = F_x
else:
# 依据 DSB 的 forward process
# $ x_{k+1} = F_k^n(x_k) + \sqrt{2 \gamma_{k+1}} \epsilon $
x_next = F_x + \
(2. * self.gammas[t]).sqrt() * torch.randn_like(x)
x_all.append(x_next.clone())
''' 为 backward model 计算并收集轨迹各状态所对应的训练 target '''
# S-DSB
if self.simplify:
if self.reparam == "TR":
# terminal 重参数化的预测目标
# 参考论文公式(20)第一式
gt_all.append(x_0)
elif self.reparam == "FR":
# Flow 重参数化的预测目标
# 参考论文公式(21)第一式
vec = (x_0 - x_next) / self.coefficient(t + 1)["coeff_1"]
gt_all.append(vec)
else:
# 当不使用重参数化时, S-DSB 的预测目标就是前一个状态
gt_all.append(x.clone())
# DSB
else:
F_x_next = self.forward_F(x, x_1, t_ts)
gt_all.append(x_next + F_x - F_x_next)
x = x_next
x_all = torch.stack([x_0] + x_all).cpu() if sample else torch.cat(x_all).cpu()
gt_all = torch.cat(gt_all).cpu()
t_all = torch.cat(t_all).cpu()
return x_all, gt_all, t_all最后，设置上参数并实例化一个 noiser 对象以备用。考虑到我们的主角是 S-DSB，因此设置 simplify = True ，至于重参数化则暂且不用，即 reparam = None， 毕竟总得先玩玩 naive 版本嘛~
training_timesteps = 16
inference_timesteps = 16
gamma_min = 1e-4
gamma_max = 1e-3
simplify = True
reparam = None
device = "cuda" if torch.cuda.is_available() else "cpu"
noiser = FlowNoiser(
device,
training_timesteps, inference_timesteps,
gamma_min, gamma_max,
simplify=simplify, reparam=reparam
)DSB 建模
对于 DSB 的建模，关键就是如何在 forward 和 backward 过程中根据给定状态和时间步去计算出下一个状态。我们可以利用神经网络，将当前状态和对应时间步喂给它，让其输出预测值，这个预测值所代表的意义则根据它在训练时所接收的 ground truth 是什么来判定(DSB, S-DSB, TR-DSB, FR-DSB 均不同)。
于是，不妨实现一个 DSB class，它继承 torch.nn.Module ，然后将 forward() 方法实现为把当前状态和时间步喂给神经网络，然后返回预测值。同时，由于需要根据 ground truth 的不同来对预测值进行后处理从而得到下一个状态，因此额外实现一个 pred_next() 方法来封装这部分操作。
另外，当 backward/forward model 在训练时，其输入状态需要由另一个 model 给出，这可以参考前面在 FlowNoiser class 中实现的 trajectory_F() 方法。由于另一个 model 为 training model 计算其输入状态属于推理过程，因此在这里我们命名为 inference() 方法，它负责计算和收集一条轨迹中所有时间步所对应的三元组：（输入状态，时间步，训练目标）。
class DSB(Module):
def __init__(
self, device,
direction: str, noiser: Module,
dim_in: int, dim_out: int,
dim_hidden: int = 128, num_layers: int = 5
):
""" DSB 建模. """
super().__init__()
self.device = device
self.noiser = noiser
self.num_timesteps = self.noiser.training_timesteps
self.simplify = self.noiser.simplify
self.reparam = self.noiser.reparam
self.noiser.prepare_gammas()
self.gammas = self.noiser.gammas
# 指示是 backward or forward model
self.direction = direction
if direction not in ('b', 'f'):
raise ValueError(f"'direction' must be 'b' or 'f', got: {direction}")
# DSB 所使用的神经网络
self.network = ResMLP(
dim_in, dim_out,
dim_hidden, num_layers,
n_cond=self.num_timesteps
).to(device)
def forward(self, x_t: Tensor, t: int) -> Tensor:
""" backward/forward model 在对应时间步所进行预测. """
timestep = self.noiser.timestep_map[t]
x = self.network(x_t, timestep)
return x
def pred_next(self, x: Tensor, t: Union[Tensor, int]) -> Tensor:
""" 根据当前状态和时间步预测下一个状态(推理过程).
对于 backward model 是: x_{k+1} -> x_k, 从 x_1 开始;
对于 forward model 则是: x_k -> x_{k+1}, 从 x_0 开始. """
if self.direction == "b":
dt = -1
# backward model, 传进来的 t 是从 num_timesteps(=N) - 1 开始,
# 而状态是从 x_1 开始, 其对应的 \gamma 系数应该是 \gamma_N
# 因此要 +1
coeff_t = t + 1
else:
dt = 1
coeff_t = t
# 当前状态和下一状态的 \bar{\gamma} & 1 - \bar{\gamma} 系数
coeff_0, coeff_1 = align_shape(x, self.noiser.coefficient(coeff_t)).values()
coeff_0_next, coeff_1_next = align_shape(x, self.noiser.coefficient(coeff_t + dt)).values()
# 对于 backward model, 是 x_{k}；
# 对于 forward model, 是 x_{k+1}.
pred = self.forward(x, t)
if self.reparam == "TR":
if self.direction == "b":
# bakcward model 预测的终点是 x_0
x_0 = pred
# 根据线性插值的原理由预测的终点和当前状态计算出起点
x_1 = (x - coeff_0 * x_0) / coeff_1
else:
# forward model 预测的终点是 x_1
x_1 = pred
# 根据线性插值的原理由预测的终点和当前状态计算出起点
x_0 = (x - coeff_1 * x_1) / coeff_0
# 根据线性插值的原理由两个端点计算出下一状态
x_next = coeff_0_next * x_0 + coeff_1_next * x_1
elif self.reparam == "FR":
# flow 重参数化情况下, 模型预测的是由当前指向终点的向量.
vec = pred
if self.direction == "b":
# 根據论文公式(19)第一式计算 backward 下一状态的均值
x_next = x + (coeff_0_next - coeff_0) * vec
else:
# 根據论文公式(19)第二式计算 forward 下一状态的均值
x_next = x + (coeff_1_next - coeff_1) * vec
else:
x_next = pred
return x_next
@torch.no_grad()
def inference(self, x: Tensor, sample: bool = False):
""" 记录推理过程的完整轨迹, 同时也可作为模型训练的 target.
若是 backward model 的推理, 则记录由 x_1 至 x_0 的各个状态, 同时
为 forward model 的训练计算 target;
同理, 若是 forward model 的推理, 则记录由 x_0 至 x_1 各个状态,
同时为 backward model 的训练计算 target."""
ones = torch.ones((x.size(0),), dtype=torch.long, device=self.device)
# 当前 model 推理的起点, 也是另一个 model 的终点
x_raw = x.clone()
x_all, gt_all, t_all = [], [], []
for t in range(self.num_timesteps):
''' 收集各时间步 '''
t_ts = ones * t
# 若是 backward model 的推理过程,
# 则方向为 x_1 -> x_0, 时间步由 N-1 至 0
if self.direction == 'backward':
t_ts = self.num_timesteps - 1 - t_ts
t_all.append(t_ts)
''' 收集各个状态 '''
# 若是 backward model, 则是 x_{k+1} -> x_k;
# 若是 forward model, 则是 x_k -> x_{k+1}.
mean_x_next = self.pred_next(x, t_ts)
# 若仅仅是采样过程且到了最后一步, 则不加噪声项
if sample and t == self.num_timesteps - 1:
x_next = mean_x_next
else:
# 由于 noise scale 呈对称形式, 因此尽管是 backward model,
# 本来 gamma 的索引应该是 num_timesteps - 1 - t, 这里直接用 t 代替也无影响
x_next = mean_x_next + \
(2. * self.gammas[t]).sqrt() * torch.randn_like(x)
x_all.append(x_next.clone())
''' 为另一个 model 计算并收集 target. '''
# S-DSB
if self.simplify:
if self.reparam == "TR":
# 论文公式(20)
gt_all.append(x_raw)
elif self.reparam == "FR":
# 论文公式(21)
gt_all.append((x_raw - x_next) / ((t + 1) / self.num_timesteps))
else:
# 论文公式(14)
gt_all.append(x.clone())
# DSB
else:
mean_mean_x_next = self.pred_next(x_next, t_ts)
# 论文公式(10)
gt_all.append(x_next + mean_x_next - mean_mean_x_next)
x = x_next.clone()
x_all = torch.stack([x_raw] + x_all).cpu() if sample else torch.cat(x_all).cpu()
gt_all = torch.cat(gt_all).cpu()
t_all = torch.cat(t_all).cpu()
return x_all, gt_all, t_all需要注意的是，由于 backward model 在推理(采样过程)时，是由 x_1 开始逐步过渡到 x_0 的，因此时间方向是 num_timesteps \- 1 ~ 0，而 inference() 方法的 for loop 中，时间是 0 ~ num_timesteps \- 1 的， 于是要判断一下：若当前 model 是 backward model，则对时间变量 t 进行“翻转”，即让其从 num_timesteps \- 1 开始递减至 0， 这个逻辑对应以上代码 t_ts = self.num_timesteps \-1- t_ts 。
另外提醒一下下，这样实现后，forward model 会从时间步 0 开始进行采样，直至最后的时间步 num_timesteps \- 1 时采样出对应 x_1 的结果；而 backward model 在采样生成时则由时间步 num_timesteps \- 1 开始，递减至时间步 0 时采样出对应 x_0 的结果，两者是某种 dual process。
神经网络的实现
上一节没有给出 DSB 所用的神经网络的实现，即上面的 self.network，现在补上。
以下选用一种比较简单的网络结构，它由一系列 MLP block 和 ResNet block 组成，前者是全连接层接激活函数，后者是在前者的基础上加入了残差连接，同时还使用了 AdaLayerNorm 作为归一化层，它将 timestep embedding 作为 scale 和 shift 因子，施加于经过 LayerNorm 作用后的特征上。
class AdaLayerNorm(Module):
r"""
Norm layer modified to incorporate timestep embeddings.
Parameters:
embedding_dim (`int`): The size of each embedding vector.
num_embeddings (`int`): The size of the embeddings dictionary.
"""
def __init__(self, num_embeddings: int, embedding_dim: int):
super().__init__()
self.emb = nn.Embedding(num_embeddings, embedding_dim)
self.silu = nn.SiLU()
self.linear = nn.Linear(embedding_dim, embedding_dim * 2)
self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False)
def forward(self, x: Tensor, timestep: Tensor) -> torch.Tensor:
x = self.norm(x)
emb = self.linear(self.silu(self.emb(timestep)))
scale, shift = torch.chunk(emb, 2, -1)
x_embed = x * scale + shift
return x + x_embed
class ResBlock(Module):
def __init__(self, dim_in, dim_out, bias=True, n_cond=1000):
super().__init__()
self.dim_in = dim_in
self.dim_out = dim_out
self.dense = nn.Linear(self.dim_in, self.dim_out, bias=bias)
self.n_cond = n_cond
if n_cond > 0:
self.norm = AdaLayerNorm(n_cond, self.dim_out)
else:
self.norm = nn.LayerNorm(dim_out)
self.activation = nn.SiLU(inplace=True)
if self.dim_in != self.dim_out:
self.skip = nn.Linear(self.dim_in, self.dim_out, bias=False)
else:
self.skip = None
def forward(self, x, t):
identity = x
if self.skip is not None:
identity = self.skip(identity)
x = self.dense(x)
norm_inputs = (x, t) if self.n_cond > 0 else (x,)
x = self.norm(*norm_inputs)
x += identity
x = self.activation(x)
return x
class BasicBlock(Module):
def __init__(self, dim_in, dim_out, bias=True):
super().__init__()
self.dim_in = dim_in
self.dim_out = dim_out
self.dense = nn.Linear(self.dim_in, self.dim_out, bias=bias)
self.activation = nn.ReLU(inplace=True)
def forward(self, x: Tensor) -> Tensor:
out = self.dense(x)
out = self.activation(out)
return out
class ResMLP(Module):
def __init__(self, dim_in, dim_out, dim_hidden, num_layers, bias=True, n_cond=1000):
super().__init__()
self.dim_in = dim_in
self.dim_out = dim_out
self.dim_hidden = dim_hidden
self.num_layers = num_layers
net = []
for l in range(num_layers):
if l == 0:
net.append(BasicBlock(self.dim_in, self.dim_hidden, bias=bias))
elif l != num_layers - 1:
net.append(ResBlock(self.dim_hidden, self.dim_hidden,
bias=bias, n_cond=n_cond))
else:
net.append(nn.Linear(self.dim_hidden, self.dim_out, bias=bias))
self.net = nn.ModuleList(net)
def forward(self, x, t):
for l in range(self.num_layers):
net_inputs = (x, t) if l not in (0, self.num_layers - 1) else (x,)
x = self.net[l](*net_inputs)
return x数据集
数据集我们就用作者在其中一个实验中所用到的棋盘格(checkerboard)数据和旋转风车(pinwheel)数据，并将前者作为 p_{data}p_{data} ，后者作为 p_{prior}p_{prior} ：
将两个数据集实现为两个类，并依照惯例继承 torch.utils.data.Dataset class：
import math
import numpy as np
from torch.utils.data import Dataset
class Checkerboard(Dataset):
def __init__(self, size=8, grid_size=4):
self.size = size
self.grid_size = grid_size
self.checkboard = torch.tensor([[i, j] for i in range(grid_size) for j in range(grid_size) if (i + j) % 2 == 0])
grid_pos = torch.randint(low=0, high=self.checkboard.shape[0], size=(self.size,), dtype=torch.int64)
self.data = torch.rand(size=(self.size, 2), dtype=torch.float32) + self.checkboard[grid_pos].float()
self.data = self.data / self.grid_size * 2 - 1
def __len__(self):
return self.size
def __getitem__(self, idx):
return self.data[idx]
class Pinwheel(Dataset):
def __init__(self, npar: int):
self.size = npar
radial_std = 0.3
tangential_std = 0.1
num_classes = 7
num_per_class = math.ceil(npar / num_classes)
rate = 0.25
rads = np.linspace(0, 2 * np.pi, num_classes, endpoint=False)
features = np.random.randn(num_classes*num_per_class, 2) \
* np.array([radial_std, tangential_std])
features[:, 0] += 1.
labels = np.repeat(np.arange(num_classes), num_per_class)
angles = rads[labels] + rate * np.exp(features[:, 0])
rotations = np.stack([np.cos(angles), -np.sin(angles), np.sin(angles), np.cos(angles)])
rotations = np.reshape(rotations.T, (-1, 2, 2))
x = .4 * np.random.permutation(np.einsum("ti,tij->tj", features, rotations))
self.init_sample = torch.from_numpy(x).float()
def __len__(self):
return self.size
def __getitem__(self, idx: int) -> Tensor:
return self.init_sample[idx]
data_size = 2 ** 26
pinwheel_dataset = Pinwheel(data_size)
checkerboard_dataset = Checkerboard(size=data_size, grid_size=8)老套路，实例化 DataLoader 以便在训练中取数据(batched data)：
from torch.utils.data import DataLoader
def create_data_loader(
dataset: Dataset, batch_size: int, shuffle: bool = False,
num_workers: int = 0, pin_memory: bool = False
) -> DataLoader:
return DataLoader(
dataset, batch_size=batch_size, shuffle=shuffle,
num_workers=num_workers, pin_memory=pin_memory
)
batch_size = 2 ** 16
pinwheel_data_loader = create_data_loader(pinwheel_dataset, batch_size, num_workers=2, pin_memory=True)
checkerboard_data_loader = create_data_loader(checkerboard_dataset, batch_size, num_workers=2, pin_memory=True)为了检验数据集实现有没有毛病，额外实现一个可视化函数来看看以上实现的两个数据集长什么样子。
import matplotlib.pyplot as plt
def show_2d_data(data: Tensor):
plt.figure(figsize=(3, 3))
plt.scatter(data[:, 0], data[:, 1])
plt.xlim(-1.1, 1.1)
plt.ylim(-1.1, 1.1)
plt.show()
plt.close()分别从两个数据集所对应的 data loader 中各自取出1个 batch 的数据进行观察：
pinwheel_batch = next(iter(pinwheel_data_loader))
checkerboard_batch = next(iter(checkerboard_data_loader))观察 checkerboard 数据是否真的长成棋盘格的样子：
show_2d_data(checkerboard_batch)
嗯，看效果还行。接着再看看 pinwheel 数据是否真的如旋转风车一般帅气：
show_2d_data(pinwheel_batch)
咦，也还可以，虽然有瑕疵，但我们大人有大量，不吹毛求疵~
训练 pipeline
一切准备就绪，如今可以来实现训练 pipeline 了。我们使用一个 Runner class 来封装，其核心部分包括：
从两个数据集的 data loader 中取出成对数据：(
x_0,x_1)为 training model 取出 batched 三元组：(输入状态，时间步，训练目标)
train loop: 计算 loss 并更新模型参数
evaluation: 每隔一段时间观察模型的采样效果
import os
import time
import tqdm
import matplotlib.animation as anime
from torch.optim import AdamW
class Runner:
def __init__(
self, device,
data_loader_0: DataLoader, data_loader_1: DataLoader,
noiser: Module, backward_model: Module, forward_model: Module,
lr: float = 1e-3, weight_decay: float = 0.,
save_path: str = '.'
):
self.device = device
self.save_path = save_path
self.data_loader_0, self.data_loader_1 = data_loader_0, data_loader_1
self.data_iter_0, self.data_iter_1 = iter(data_loader_0), iter(data_loader_1)
self.num_batches = len(data_loader_0)
self.batch_size = data_loader_0.batch_size
# 每一对数据 (x_0, x_1) 都对应有完整的一条轨迹(x_{k-1}, x_k, x_{k+1}..),
# 因此先缓存下一批数据避免每次迭代都重新计算
self.cache_size = self.cnt = self.num_batches * self.batch_size * 4
self.loss_fn = nn.MSELoss()
self.noiser = noiser
self.backward_model = backward_model.to(device)
self.forward_model = forward_model.to(device)
self.backward_optimizer = AdamW(
self.backward_model.parameters(),
lr=lr,
weight_decay=weight_decay
)
self.forward_optimizer = AdamW(
self.forward_model.parameters(),
lr=lr,
weight_decay=weight_decay
)
self.model_dict = {
'backward': self.backward_model,
'forward': self.forward_model
}
self.optimizer_dict = {
'backward': self.backward_optimizer,
'forward': self.forward_optimizer
}
self.direction = 'backward'
def get_paired_data(self) -> Tuple[Tensor]:
try:
x_0, x_1 = next(self.data_iter_0), next(self.data_iter_1)
except StopIteration:
self.data_iter_0 = iter(self.data_loader_0)
self.data_iter_1 = iter(self.data_loader_1)
x_0, x_1 = next(self.data_iter_0), next(self.data_iter_1)
return x_0.to(self.device), x_1.to(self.device)
def get_batch(self, epoch: int) -> Tuple[Tensor]:
# 当缓存里的数据已经取完, 则重新计算
if self.cnt + self.batch_size > self.cache_size:
self.x_cache, self.gt_cache, self.t_cache = [], [], []
num_pairs = math.ceil(self.cache_size / (self.batch_size * self.noiser.num_timesteps))
pbar = tqdm.trange(num_pairs, desc=f"Caching data on epoch {epoch} for {self.direction} model..")
for _ in pbar:
x_0, x_1 = self.get_paired_data()
with torch.no_grad():
# 首个 epoch 训练 backward model
if epoch == 0:
if self.direction != 'backward':
raise RuntimeError("Epoch 0 should be backward model.")
# 由于 forward model 还未训练, 因此只能根据线性插值计算 backward model 的 data & target.
x_all, gt_all, t_all = self.noiser.trajectory_F(x_0, x_1)
else:
# 若当前训练的是 forward model, 则其 data & target 由 backward model 推理生成;
# 反之, 若训练的是 backward model, 则 data & target 由 forward model 推理生成.
model = self.model_dict['backward' if self.direction == 'forward' else 'forward'].eval()
x_start = x_1 if self.direction == 'forward' else x_0
x_all, gt_all, t_all = model.inference(x_start)
self.x_cache.append(x_all)
self.gt_cache.append(gt_all)
self.t_cache.append(t_all)
self.x_cache = torch.cat(self.x_cache).cpu()
self.gt_cache = torch.cat(self.gt_cache).cpu()
self.t_cache = torch.cat(self.t_cache).cpu()
self.cnt = 0
self.cache_indices = torch.randperm(self.x_cache.size(0))
# 每次取1个 batch 并记录取到哪里
indices = self.cache_indices[self.cnt:self.cnt + self.batch_size]
self.cnt += self.batch_size
x_batch = self.x_cache[indices].to(self.device)
gt_batch = self.gt_cache[indices].to(self.device)
t_batch = self.t_cache[indices].to(self.device)
return x_batch, gt_batch, t_batch
def train(
self, n_epochs: int, repeat_per_epoch: int,
log_interval: int = 128, eval_interval: int = 1024
):
steps_per_epoch = self.num_batches * repeat_per_epoch
self.cache_size = min(self.cache_size, steps_per_epoch * self.batch_size)
for epoch in range(n_epochs):
self.noiser.train()
# 两个 model 交替训练
self.direction = 'backward' if epoch % 2 == 0 else 'forward'
model, optimizer = self.model_dict[self.direction], self.optimizer_dict[self.direction]
model.train()
optimizer.zero_grad()
self.cnt = self.cache_size
pbar = tqdm.tqdm(total=steps_per_epoch)
# 使用 ema loss 来观察整体趋势
ema_loss, ema_loss_w = None, lambda x: min(0.99, x / 10)
for step in range(steps_per_epoch):
x_t, gt, t = self.get_batch(epoch)
pred = model(x_t, t)
loss = self.loss_fn(pred, gt)
loss.backward()
loss = loss.item()
optimizer.step()
optimizer.zero_grad()
ema_loss = loss if ema_loss is None \
else (ema_loss * ema_loss_w(step) + loss * (1 - ema_loss_w(step)))
if (step + 1) % log_interval == 0 or step == steps_per_epoch - 1:
info = f'Epoch: [{epoch}]/[{n_epochs}]; Step: {step}; Direction: {self.direction}; Loss: {loss:.14f}; Ema Loss: {ema_loss:.14f}'
pbar.set_description(info, refresh=False)
pbar.update(step + 1 - pbar.n)
# 训练到一定步数就推理看看效果
if (step + 1) % eval_interval == 0:
self.evaluate(epoch, step + 1)
model.train()
self.evaluate(epoch, steps_per_epoch, last_step=True)
@torch.no_grad()
def evaluate(self, epoch: int, step: int, last_step: bool = False):
self.backward_model.eval()
self.forward_model.eval()
x_0, x_1 = self.get_paired_data()
qs = self.backward_model.inference(x_1, sample=True)[0]
# 首个 epoch forward model 还未训练, 因此插值计算出整条 forward 轨迹.
ps = self.forward_model.inference(x_0, sample=True)[0] \
if epoch else self.noiser.trajectory_F(x_0, x_1, sample=True)[0]
# 画出两个 model 最后1步生成的样本
self.draw_sample(qs[-1], epoch, step, subfix='_0')
self.draw_sample(ps[-1], epoch, step, subfix='_1')
# 画出两个方向对应的整条轨迹
if (epoch + 1) % 2 == 0 and last_step:
self.draw_trajectory(qs, epoch, subfix='_0')
self.draw_trajectory(ps, epoch, subfix='_1')
def draw_sample(
self, sample: Tensor,
epoch: int, step: int,
xrange=(-1, 1), yrange=(-1, 1),
subfix: str = None
):
sample = sample.cpu().numpy()
save_path = os.path.join(
self.save_path,
f'sample' + (subfix if subfix is not None else ''),
)
os.makedirs(save_path, exist_ok=True)
plt.figure(figsize=(10, 10))
plt.scatter(sample[:, 0], sample[:, 1], s=1)
plt.xlim(xrange[0] - 0.1, xrange[1] + 0.1)
plt.ylim(yrange[0] - 0.1, yrange[1] + 0.1)
plt.savefig(os.path.join(save_path, f'ep{epoch}_it{step}.jpg'))
plt.close()
def draw_trajectory(self, xs: Tensor, epoch: int, xrange=(-1, 1), yrange=(-1, 1), subfix: str = None):
save_path = os.path.join(self.save_path,
f'trajectory' + (subfix if subfix is not None else ''),
f'ep{epoch}'
)
os.makedirs(save_path, exist_ok=True)
k = 1000
xs = xs.cpu().numpy()
plt.figure(figsize=(10, 10))
plt.scatter(
x=np.reshape(xs[:, :k, 0], -1),
y=np.reshape(xs[:, :k, 1], -1),
s=1, cmap='viridis', vmin=0, vmax=1,
c=np.reshape(np.repeat(np.expand_dims(np.arange(xs.shape[0]), 1), k, axis=1), -1) / xs.shape[0]
)
plt.xlim(xrange[0] - 0.1, xrange[1] + 0.1)
plt.ylim(yrange[0] - 0.1, yrange[1] + 0.1)
plt.savefig(os.path.join(save_path, 'trajectory.jpg'))
plt.close()
self.draw_animation(xs, save_path, xrange=xrange, yrange=yrange)
def draw_animation(self, xs: Tensor, save_path: str, xrange=(-1, 1), yrange=(-1, 1)):
clamp = lambda x, a, b: int(min(max(x, a), b))
st = time.perf_counter()
num_timesteps, batch_size = xs.shape[0], xs.shape[1]
steps_per_second = clamp(num_timesteps / 100, 1, 10)
frames_per_second = clamp(num_timesteps / 10, 1, 10)
num_seconds = num_timesteps / frames_per_second / steps_per_second + 3
print('plotting point cloud animation ......', end='', flush=True)
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim(xrange[0] - 0.1, xrange[1] + 0.1)
ax.set_ylim(yrange[0] - 0.1, yrange[1] + 0.1)
scatter = ax.scatter([], [], s=1, c=[], cmap='viridis', vmin=0, vmax=1)
def animate(j):
j = min((j + 1) * steps_per_second, xs.shape[0])
cc = np.arange(j) / (num_timesteps - 1)
cc = np.reshape(np.repeat(np.expand_dims(cc, axis=1), batch_size, axis=1), -1)
scatter.set_offsets(np.reshape(xs[j - 1], (-1, 2)))
scatter.set_array(cc[-batch_size:])
return scatter,
ani = anime.FuncAnimation(fig, animate, frames=int(num_seconds*frames_per_second), interval=1000/frames_per_second, repeat=False, blit=True)
try:
ani.save(os.path.join(save_path, 'trajectory.mp4'), writer=anime.FFMpegWriter(fps=frames_per_second, codec='h264'), dpi=100)
except:
ani.save(os.path.join(save_path, 'trajectory.gif'), writer=anime.PillowWriter(fps=frames_per_second), dpi=100)
plt.close(fig)
print(f' done! ({time.perf_counter() - st:.3f}s)')以上 get_batch() 方法就是用于取出 batched 三元组的，它事先将整条轨迹都 cache 下来，以便在每个训练迭代中，要取数据时就直接从其中取，避免每次迭代时在线计算。同时，可以看到其逻辑与前面所述一致：在第1个 epoch，由于 forward model 还未训练，因此三元组需要依赖 noiser 给出；在而后的每个 epoch，training model 所需的三元组则由另一个 model 的 inference() 方法给出。
另外，考虑到每轮(epoch)训练效果依赖于上一轮，于是这里的训练过程并非将 data loader 的数据取完即算1个 epoch，而是取完后再将所有数据重复(重新实例化 data loader)，这样重复搞一定次数后(我们这里设为8次)才算1个 epoch 结束。
Last step，实例化 DSB class 和 Runner class，开跑！
dim_in = dim_out = 2
backward_model = DSB(device, 'b', noiser, 2, 2)
forward_model = DSB(device, 'f', noiser, 2, 2)
lr = 1e-3 #@param {'type': 'number'}
save_path = 'exp2d' #@param {'type': 'string'}
runner = Runner(
device,
checkerboard_data_loader, pinwheel_data_loader,
noiser, backward_model, forward_model,
lr=lr, save_path=save_path
)
n_epochs = 41 #@param {'type': 'integer'}
repeat_per_epoch = 8 #@param {'type': 'integer'}
runner.train(n_epochs, repeat_per_epoch)Woohoo~ 训起来了:
Caching data on epoch 0 for backward model..: 97%|█████████▋| 248/256 [00:18<00:00, 10.64it/s]
Caching data on epoch 0 for backward model..: 98%|█████████▊| 250/256 [00:18<00:00, 11.79it/s]
Caching data on epoch 0 for backward model..: 98%|█████████▊| 252/256 [00:18<00:00, 13.35it/s]
Caching data on epoch 0 for backward model..: 99%|█████████▉| 254/256 [00:18<00:00, 14.26it/s]
Caching data on epoch 0 for backward model..: 100%|██████████| 256/256 [00:18<00:00, 13.72it/s]
Epoch: [0]/[41]; Step: 127; Direction: backward; Loss: 0.00270029623061; Ema Loss: 0.05426615317930: 2%|▏ | 128/8192 [00:23<24:32, 5.47it/s]
Epoch: [0]/[41]; Step: 255; Direction: backward; Loss: 0.00254723778926; Ema Loss: 0.01683711398279: 3%|▎ | 256/8192 [00:24<10:39, 12.41it/s]
Epoch: [0]/[41]; Step: 383; Direction: backward; Loss: 0.00242871278897; Ema Loss: 0.00645837263667: 5%|▍ | 384/8192 [00:25<06:12, 20.97it/s]
Epoch: [0]/[41]; Step: 511; Direction: backward; Loss: 0.00244712480344; Ema Loss: 0.00357424189259: 6%|▋ | 512/8192 [00:26<04:07, 31.00it/s]
Epoch: [0]/[41]; Step: 639; Direction: backward; Loss: 0.00252332701348; Ema Loss: 0.00277143061942: 8%|▊ | 640/8192 [00:27<02:59, 42.12it/s]
Epoch: [0]/[41]; Step: 767; Direction: backward; Loss: 0.00245906645432; Ema Loss: 0.00254573895120: 9%|▉ | 768/8192 [00:28<02:18, 53.75it/s]
Epoch: [0]/[41]; Step: 895; Direction: backward; Loss: 0.00243897456676; Ema Loss: 0.00248167126902: 11%|█ | 896/8192 [00:30<01:52, 65.07it/s]
Epoch: [0]/[41]; Step: 1023; Direction: backward; Loss: 0.00244102231227; Ema Loss: 0.00245851060842: 12%|█▎ | 1024/8192 [00:31<01:35, 74.72it/s]
Epoch: [1]/[41]; Step: 7167; Direction: forward; Loss: 0.00111923017539; Ema Loss: 0.00111609658618: 88%|████████▊ | 7168/8192 [02:01<00:09, 106.24it/s]
Epoch: [1]/[41]; Step: 7295; Direction: forward; Loss: 0.00110315659549; Ema Loss: 0.00110453149761: 89%|████████▉ | 7296/8192 [02:03<00:09, 96.95it/s]
Epoch: [1]/[41]; Step: 7423; Direction: forward; Loss: 0.00110766326543; Ema Loss: 0.00111066790597: 91%|█████████ | 7424/8192 [02:04<00:07, 101.86it/s]
Epoch: [1]/[41]; Step: 7551; Direction: forward; Loss: 0.00109377037734; Ema Loss: 0.00111771681568: 92%|█████████▏| 7552/8192 [02:05<00:06, 105.26it/s]
Epoch: [1]/[41]; Step: 7679; Direction: forward; Loss: 0.00109407398850; Ema Loss: 0.00110658616989: 94%|█████████▍| 7680/8192 [02:06<00:04, 108.15it/s]
Epoch: [1]/[41]; Step: 7807; Direction: forward; Loss: 0.00110605405644; Ema Loss: 0.00110584671991: 95%|█████████▌| 7808/8192 [02:07<00:03, 110.43it/s]
Epoch: [1]/[41]; Step: 7935; Direction: forward; Loss: 0.00110646407120; Ema Loss: 0.00110231996900: 97%|█████████▋| 7936/8192 [02:08<00:02, 112.67it/s]
Epoch: [1]/[41]; Step: 8063; Direction: forward; Loss: 0.00111326680053; Ema Loss: 0.00110731386783: 98%|█████████▊| 8064/8192 [02:10<00:01, 108.24it/s]
Epoch: [1]/[41]; Step: 8191; Direction: forward; Loss: 0.00108781910967; Ema Loss: 0.00110913189688: 100%|██████████| 8192/8192 [02:11<00:00, 103.22it/s]
Epoch: [1]/[41]; Step: 8191; Direction: forward; Loss: 0.00108781910967; Ema Loss: 0.00110913189688: 100%|██████████| 8192/8192 [02:12<00:00, 61.92it/s]最终，要是你运气好，就可以炼出以下效果：
运气不好嘛..就
玩后感
玩了一把 S-DSB，虽然它扬言要为 DSB 治病，但 CW 感觉它自身也是有不少毛病的，比如：
尽管简化了 loss 函数，但收敛还是慢且不容易(就算使用了重参数化技术和利用了 pre-trained SGM)
收敛慢导致无法 scale up 到更大数据集
paper 所提出的重参数化方法基于一定假设，这与实际情况存在冲突，导致在实践中可能会造成负面影响
另外，training model 的参考目标轨迹是由另一个 model 推理出来的，可能会严重偏离真实的目标轨迹。如果另一个 model 没训好，那么 training model 的 target 就是 “胡扯”，而另一个 model 是上一轮的训练成果，于是每轮的训练效果都极大依赖于上一轮的收敛效果。
更可怕的是，这会导致每一轮在训练时本需满足的 “边界条件” 都不满足(i.e. 或 )。比如在偶数轮 优化 (对应 forward 轨迹)时本需满足 ，但由于 backward model 不够优秀(上一轮没训好)、其采样生成的结果 ，也就是采样结果不服从 ，从而就会导致 。在奇数轮 优化 (对应 backward 轨迹)时同理，也容易造成 。
诶~！打住，别以为我是个怼怼，所谓欲扬先抑嘛，这不，我接下来就要好好夸夸这孩子呢！
对比起 SGM(扩散模型)，这种 DSB 框架的天花板在先天上是更高的，毕竟它没有强制要求 是简单的高斯分布甚至都没有 的概念！相互转换的两个分布都可以是有意义的任意分布，这同时也为条件生成带来了优秀的基因。同时，比起 SGM 只能 单向生成， DSB 是天然就能够 双向生成的，这可就好玩多了，妥妥的不无聊的风格!
Anyway, 在主流之外就是要有一些不一样的东西出现才好，尽管有利有弊，但在整体上是能促进历史的螺旋式前进与发展的。
投稿作者为『自动驾驶之心知识星球』特邀嘉宾，欢迎加入交流！
① 全网独家视频课程
BEV感知、BEV模型部署、BEV目标跟踪、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、车道线检测、轨迹预测、在线高精地图、世界模型、点云3D目标检测、目标跟踪、Occupancy、cuda与TensorRT模型部署、大模型与自动驾驶、Nerf、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频（扫码即可学习）
② 国内首个自动驾驶学习社区
国内最大最专业，近3000人的交流社区，已得到大多数自动驾驶公司的认可！涉及30+自动驾驶技术栈学习路线，从0到一带你入门自动驾驶感知（2D/3D检测、语义分割、车道线、BEV感知、Occupancy、多传感器融合、多传感器标定、目标跟踪）、自动驾驶定位建图（SLAM、高精地图、局部在线地图）、自动驾驶规划控制/轨迹预测等领域技术方案、大模型、端到端等，更有行业动态和岗位发布！欢迎扫描下方二维码，加入自动驾驶之心知识星球，这是一个真正有干货的地方，与领域大佬交流入门、学习、工作、跳槽上的各类难题，日常分享论文+代码+视频
③【自动驾驶之心】技术交流群
自动驾驶之心是首个自动驾驶开发者社区，聚焦感知、定位、融合、规控、标定、端到端、仿真、产品经理、自动驾驶开发、自动标注与数据闭环多个方向，目前近60+技术交流群，欢迎加入！扫码添加汽车人助理微信邀请入群，备注：学校/公司+方向+昵称（快速入群方式）
④【自动驾驶之心】全平台矩阵
5756
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
