资讯详情

PWC-Net光流训练实战:多尺度监督、NaN排查与混合精度调优

📅 2026/10/27 5:04:23 | 华诺云谱 👁 阅读
PWC-Net光流训练实战:多尺度监督、NaN排查与混合精度调优
1. PWC-Net训练回路里究竟在算什么把网络结构拼出来、跑通一次前向、打印出六个尺度的光流张量这只是复现PWC-Net的上半场。真正让人卡住三周甚至更久的是下半场——把训练循环跑起来让loss稳定下降让EPE指标在验证集上真的往好的方向走。我在Paddle上重写这套光流估计训练流程时第一次跑出来的日志是这样的前200个iteration loss从4.7掉到1.9看起来很美然后第600个iteration直接跳到nan之后无论怎么重启都在同一个位置爆掉。后面翻了整整两天的中间张量才定位到原因这段经历后面会详细讲。先说清楚训练阶段在做什么。PWC-Net本质上是一个金字塔式的由粗到细光流估计网络输入是两张连续帧img1、img2输出是六个尺度上的光流场。训练时这六个输出每一层都要参与监督最粗的一层负责捕捉大位移最细的一层负责修正细节整个梯度要同时穿过特征金字塔、warping采样层、cost volume相关层和光流估计头。这意味着反向传播的计算图比前向推理长得多显存峰值通常是单纯前向的三到四倍这也是为什么很多人结构复现成功、训练却一直OOM或者被迫把batch压到1。对第一次接触光流训练的读者我需要先把这个任务和普通的图像分类、分割区分开。分类任务的监督信号是一个标量标签分割任务是一张和原图同尺寸的mask而光流的监督信号是每个像素的二维位移向量并且这个向量的量纲会随着金字塔层数变化——这是后面大量坑的根源。第二帧图像是通过第一帧加上光流采样得到的这个采样操作在网络内部和损失函数里都会出现一旦坐标归一化方式或者align_corners设置和预处理对不上损失会以一种非常隐蔽的方式失真。这一篇要讲的就是训练部分的完整落地数据管线怎么搭、多尺度损失怎么写、优化器和学习率怎么配、日志里那些异常怎么排查、验证集怎么评估、checkpoint怎么管。适合已经能把网络跑通、准备正式开始训练的读者也适合正在被loss不降或者NaN折磨的人直接对照排查。1.1 从“能前向”到“能训练”之间隔着什么前向推理时你只需要关心输出对不对中间的任何张量用完就可以丢。训练时不一样autograd需要保留所有参与梯度计算的中间结果六个尺度的特征图、每一层的warp结果、每一层的cost volume、每一层的初步光流和refine之后的最终光流。粗略估算一下在384×448的输入下六层金字塔每层都要保留一对特征图通道数从16到128递增和对应的cost volume通道数是search range的两倍加一显存占用很容易到几个GB。我在一张中等显存的卡上实测batch size为4、输入384×448、不开混合精度时峰值显存大概是11GB左右把batch降到2才能勉强跑起来但batch太小又会让BatchNorm的统计量不稳定。这里有个容易被忽略的点PWC-Net的特征提取部分在两层之间是共享权重的吗论文里的设计是每一层金字塔各有一套独立的卷积不做跨层共享。这意味着参数量是随层数线性增长的而不是共享后的一份。有人在复现时为了省参数改成共享结果收敛速度和最终精度都不对。我建议严格按照论文的独立卷积来实现参数量大概在几百万这个量级对显存的影响主要来自激活值而不是权重。还有一个隐藏成本是数据加载。光流训练集通常是成对的图片加上一个流文件FlyingChairs大约有两万多对样本每对包含两张384×512的图和一个约1.5MB的flo文件。如果用Python原生循环去读单线程读一张图加一个flo文件差不多要20到30毫秒GPU跑一个iteration只要不到100毫秒数据管线立刻成为瓶颈。这也是为什么训练PWC-Net时DataLoader的worker数量和prefetch策略必须认真调不然后面所有的调参都是在和IO较劲。1.2 训练阶段的数据分布为什么不能一个数据集跑到底PWC-Net的原始训练流程是分三段走的这一点在复现时经常被简化成“拿一个数据集从头训到尾”然后发现指标怎么都上不去。第一段用合成数据集类似FlyingChairs这种由真实图片合成位移的集合做大规模预训练样本量在二十万这个量级目的是让网络学会基本的匹配和位移估计能力。第二段换到更复杂的三维渲染合成数据集做微调这类数据包含更丰富的遮挡、光照变化和大位移。第三段才转到真实采集的数据集上做小学习率微调样本量通常只有几百到几千对。这个分段的逻辑是合成数据的位移是精确已知的可以作为强监督真实数据集的标注稀疏或者存在无效区域只能小学习率微调否则容易过拟合到少量样本上。我在复现时试过直接用真实数据集从头训结果训练集loss降得飞快验证集EPE卡在一个很差的水平不动典型的过拟合。后来老老实实按三段走同样的网络结构最终指标提升了接近三成。训练阶段数据规模量级学习率量级迭代次数主要作用合成预训练二十万对1e-430万学习基础匹配与位移估计三维渲染微调两万对1e-410万适应遮挡与大位移真实数据微调千级1e-510万贴合真实场景分布表里的迭代次数是按batch size为8折算的如果你把batch调大迭代次数要相应减少保持总样本吞吐量大致相当。这个换算关系很多人在复现时会忘记导致实际训练量翻倍或者减半指标自然对不上。2. 六层金字塔的监督信号怎么组织才不会互相打架多尺度监督是PWC-Net训练部分最核心也最容易写错的地方。网络在每一层金字塔上都输出一个光流预测这六个预测需要和对应尺度上的真值光流做比较。听起来简单但“对应尺度”这四个字里藏着一堆细节真值光流怎么缩放到该尺度、缩放时位移值要不要乘尺度因子、无效区域怎么mask、每层损失权重怎么分配。这四件事任何一件写错训练都不会崩但指标就是上不去属于典型的“沉默错误”。2.1 为什么每一层都要单独监督而不是只监督最后一层如果只监督最细的一层梯度要穿过五层warp和五层cost volume才能传到最粗层衰减非常严重最粗层基本学不到东西而最粗层恰恰负责大位移估计。PWC-Net的论文设计里最粗层先估计一个粗略光流然后逐层warp上一层的图像并用cost volume做残差修正。如果最粗层没学好后面所有层的输入都是错的误差会层层放大。所以每一层都必须有直接的监督信号让粗层也能得到足够强的梯度。损失权重通常是按指数衰减的。论文里给出的一组参考值是最细层权重最大往粗层方向逐级减半左右大致落在0.32到0.005这个区间。这个分配的逻辑是最细层的光流是最终输出直接决定指标权重必须最大粗层的作用是提供一个好的初值不需要它精确到像素级权重给小一些可以避免粗层的噪声干扰细层的学习。我在实际训练里调整过这组权重。把最粗层的权重从0.005提到0.02之后合成预训练阶段的收敛速度明显变快验证集EPE在前五万次迭代内下降得更陡但训练到后期细层精度略有下降。反过来把最粗层权重降到0.001前期收敛变慢后期细层指标反而更好一点。这说明权重不是死的取决于你的训练轮次和数据分布。我的建议是先用论文的参考值跑通把这条基线记住再去动权重。2.2 尺度对齐那个漏乘的2.0让我的loss卡了三天每一层输出的光流是在该层特征图的坐标系下定义的。第0层是原图尺度第1层是原图的一半第2层是四分之一以此类推。当你把第k层的光流上采样回原图尺寸去计算损失时除了做空间上的双线性插值还必须把位移值乘以2的k次方。原因是在第k层里“位移一个像素”对应到原图上是“位移2^k个像素”。我第一次写这段代码时只做了空间上采样忘了乘尺度因子。现象是训练集loss很快降到一个很低的水平验证集EPE却始终在十几像素的量级徘徊。排查了很久才意识到粗层的预测被当成原图尺度的位移去监督了网络学到的是一个被压缩了2^k倍的位移场而损失函数因为真值也被同样错误地处理我当时真值缩放也写错了所以两边一致地错了loss看起来很正常。正确的做法有两种选一种并且全程保持一致。第一种是网络输出什么尺度就在什么尺度上算损失真值光流下采样到对应尺度并除以2的k次方第二种是把每层输出都上采样回原图尺度并乘以2的k次方真值保持原尺度。两种方式在数学上等价但第二种实现更简单也更容易和论文的评估方式对齐。我后来统一用第二种代码结构清晰很多。import paddle import paddle.nn.functional as F def multi_scale_loss(pred_flows, gt_flow, weights(0.32, 0.08, 0.02, 0.01, 0.005, 0.0025)): pred_flows: list从细到粗共6个尺度的预测 gt_flow: [B, 2, H, W] 原图尺度的真值光流 total 0.0 for level, (pred, w) in enumerate(zip(pred_flows, weights)): scale 2 ** level if scale 1: # 空间上采样 位移值乘尺度因子 pred F.interpolate( pred, sizegt_flow.shape[-2:], modebilinear, align_cornersTrue ) pred pred * float(scale) diff paddle.abs(pred - gt_flow) # 按通道求L1再对空间求均值 total total w * diff.mean() return total这段代码里有两个点需要特别注意。align_cornersTrue必须和warping层的采样设置保持一致否则上采样的插值方式会带来亚像素级的系统性偏移。另外pred * float(scale)这一步不能写在插值之前因为插值是线性的乘一个常数放在前后理论上等价但放在后面能避免中间结果的数值范围过大导致fp16溢出。还有一个细节是无效区域的mask。真实数据集的真值光流里存在大量无效像素比如被遮挡的区域或者超出图像边界的位移这些位置的流值通常是0或者一个特殊标记值。如果不做mask网络会被迫去拟合这些无意义的零位移损失会被稀释。我的做法是在读取数据时额外生成一个有效区域mask损失函数里对diff乘以mask后再求均值分母用mask的和而不是空间尺寸。def masked_l1(pred, gt, maskNone, eps1e-6): diff paddle.abs(pred - gt).sum(axis1) # [B, H, W] if mask is not None: diff diff * mask return diff.sum() / (mask.sum() eps) return diff.mean()这套写法在我后续所有实验里都复用了比直接mean稳很多尤其是在混合了不同数据集的情况下。3. 数据管线让GPU别再空等CPU光流任务的数据管线有两个特殊要求一是必须成对读取两张图并保证它们在空间上严格对齐二是必须有和图片同步的几何变换。这两点决定了你不能用现成的图像分类Dataset直接套必须自己从零写一个。3.1 Dataset的实现要点与flo文件解析标准的flo文件格式是以一个4字节的魔数开头紧接着两个int32表示宽和高然后是宽乘高乘2个float32按行优先排列顺序是(u, v)。读的时候用numpy的fromfile直接从二进制缓冲区解析比用struct逐个解要快一个数量级。import numpy as np def read_flo(path): with open(path, rb) as f: magic np.fromfile(f, np.float32, count1) if magic ! 202021.25: raise ValueError(fbad flo magic: {magic}) w np.fromfile(f, np.int32, count1)[0] h np.fromfile(f, np.int32, count1)[0] data np.fromfile(f, np.float32, count2 * w * h) flow np.resize(data, (h, w, 2)) return flow # [H, W, 2], (u, v)这里有个小坑np.fromfile读魔数时用的是float32而文件里魔数确实是按float32存的。有些版本的实现会按int32读然后比较20202125也能对上但如果你照抄了某个博客的代码发现读取报错先检查这一行的类型。Dataset类需要返回三样东西img1、img2、flow以及可选的valid mask。图片统一转成RGB三通道、float32、归一化到[0,1]区间。这一点必须和网络的输入预处理严格一致我在复现时就遇到过一处不一致Dataset里做了mean/std归一化网络里又归一化了一次等于做了两次导致输入分布完全偏移训练loss从一开始就异常高。DataLoader的配置上num_workers建议设成CPU核数的四分之一到一半use_shared_memoryTrue能显著减少进程间的数据拷贝开销。Paddle的DataLoader默认会做collate把list拼成batch只要每个样本的shape一致就不需要自定义collate_fn。3.2 几何增强必须作用在图和流上且变换方式不同这是数据管线里最容易写错的地方。随机裁剪、水平翻转、缩放这三种增强作用在图片上是像素级的重采样作用在光流上的规则完全不同。随机裁剪相对简单从同一组随机坐标里同时裁img1、img2、flow和mask保证四者空间对齐。裁剪尺寸必须是64的倍数因为六层金字塔每层下采样2倍最粗层是原图的1/32如果输入不能被64整除金字塔各层的尺寸会出现不一致的取整warp和cost volume的维度对不上直接报错。我一般用384×448或者512×384这一档。水平翻转需要特别处理。图片翻转之后光流的u分量要取反并且如果翻转是带偏移的比如中心翻转u还要做位置平移。最稳妥的写法是把flow的u通道在翻转后取负v通道保持不变。def flip_pair(img1, img2, flow): img1 img1[:, ::-1] img2 img2[:, ::-1] flow flow[:, ::-1].copy() flow[..., 0] -flow[..., 0] # u 分量取反 return img1, img2, flow缩放的处理是flow乘以缩放比例。如果你的增强流程里同时有裁剪和缩放顺序很重要先缩放再裁剪flow的数值调整只需要在缩放那一步做一次乘比例裁剪不影响流值。颜色相关的增强亮度、对比度、饱和度轻微扰动只作用在图片上flow不动。这类增强在光流任务里要克制使用因为两帧之间的光照一致性本身是网络判断对应关系的重要线索扰动过强会让网络学到错误的匹配策略。我一般把亮度扰动的幅度控制在正负0.1以内。增强操作作用对象光流调整方式常用幅度随机裁剪图/流/mask 同步无需调整384×448 或 512×384水平翻转图/流 同步u 取反概率 0.5缩放图/流 同步流值乘缩放比0.9 到 1.1亮度扰动仅图片不影响正负 0.1高斯噪声仅图片不影响极轻慎用4. 优化器、学习率调度与梯度裁剪的联动训练PWC-Net时优化器配置看着简单实际有一堆需要对齐的细节。论文里用的是Adam初始学习率1e-4在预设的迭代节点上减半。听起来很标准但因为光流网络里有权值共享的卷积和多个尺度输出梯度量级在不同层之间差异很大如果完全照搬论文而不加任何保护很容易在几千次迭代内就出现梯度爆炸。4.1 Adam的参数选择与权重衰减的实际作用Adam在Paddle里的调用很直接但有几个参数值得讨论。beta1用默认的0.9、beta2用0.999是常规选择我没有找到调整它们的强理由。真正影响训练稳定性的是权重衰减和epsilon。光流网络在小数据集微调阶段很容易过拟合加一个较小的权重衰减量级1e-4能明显改善验证集指标。epsilon建议保持默认的1e-8如果调大会在梯度很小的层上引入额外偏置。import paddle from paddle.optimizer import Adam optimizer Adam( learning_rate1e-4, parametersmodel.parameters(), beta10.9, beta20.999, epsilon1e-8, weight_decay1e-4, grad_clippaddle.nn.ClipGradByGlobalNorm(clip_norm1.0), )grad_clip这行是重点。光流网络里的warp操作本质是双线性采样它的反向传播会对采样点坐标求梯度当采样点落在图像边界附近时插值权重的导数会出现尖峰单个batch里只要有一个这样的样本全局梯度范数就可能瞬间放大几十倍。我实测过不加裁剪的版本在合成数据集上跑到两三千次迭代必爆加了ClipGradByGlobalNorm(1.0)之后能稳定跑完全程。这个裁剪阈值不需要很精细0.5到5之间都能用1.0是个比较稳的中间值。4.2 分段学习率的具体实现与常见误用论文里的调度是阶梯式的在30万次迭代的训练中分别在10万、15万、20万、25万次时把学习率减半。Paddle里可以用PiecewiseDecay一次性定义好整个阶梯。from paddle.optimizer.lr import PiecewiseDecay scheduler PiecewiseDecay( boundaries[100000, 150000, 200000, 250000], values[1e-4, 5e-5, 2.5e-5, 1.25e-5, 6.25e-6], )这里有两个常见的误用。第一个是忘记在每个iteration之后调用scheduler.step()导致学习率一直停在第一个值训练后期无法精细收敛。第二个是把scheduler传给optimizer之后又手动调用了optimizer.set_lr()两套机制打架实际学习率不可预测。我的建议是只保留一种方式把scheduler传给optimizer的learning_rate参数然后在每个iteration末尾调用一次optimizer.step()Paddle会自动驱动scheduler。微调阶段的学习率要重新初始化不能沿用上一阶段的scheduler。从合成数据转到真实数据时我把学习率直接从1e-4降到1e-5并且重新定义了一个新的阶梯调度。如果沿用旧scheduler它会认为已经走了几万步直接给你一个很低的学习率微调根本不动。5. 训练日志里那些异常信号的排查链路训练能不能跑通是一回事跑得对不对是另一回事。我在复现过程中遇到过三次比较严重的异常这里把完整的排查过程写出来方便你遇到类似现象时按图索骥。5.1 loss在一个固定迭代位置变成NaN的定位过程回到开头提到的那个场景。现象是前600次迭代正常loss从4.7平稳降到2.1左右第601次iteration突然变成nan之后从最近一个checkpoint恢复依然在同一个位置爆掉。这个“固定位置”很关键说明不是随机性的数值溢出而是某一条特定的训练样本触发了问题。我的排查顺序是这样的。第一步把batch size临时改成1遍历数据集看是哪一对样本触发的。结果定位到一对输入里第二张图的右下角有一小块纯黑区域对应的光流值是个很大的数。第二步检查数据读取发现这个flo文件本身的宽高字段和实际文件大小不匹配读取时发生了错位导致流值里混进了文件头的字节出现了几万量级的异常值。第三步确认修复方式。在Dataset里加一层防护读完之后检查flow的绝对值是否超过一个阈值比如1000像素超过就视为无效把对应区域标记为mask为0同时打印一条警告。重新跑训练NaN消失。但这个案例暴露的更深层问题是不能完全信任数据集的完整性。我在后续的实验里统一加了三个检查点写成了一个数据体检脚本训练开始前先跑一遍。检查项判据处理方式图片尺寸一致性img1/img2/flow 三者宽高相同不一致则跳过或裁剪到最小公共区域流值范围绝对值超过 1000 像素标记为无效像素mask 置 0无效值检测存在 NaN 或 Inf整对样本丢弃并记录文件头校验flo 魔数等于 202021.25校验失败直接抛异常宽高与文件大小匹配8 wh2*4 等于实际字节数不匹配则视为损坏文件这套检查跑一遍两万对样本大概几分钟能省掉后面几十个小时的无效训练。5.2 学习率正常但loss不降从warp采样参数入手第二次异常更隐蔽loss能降但降得非常慢训练到五万次迭代才到别人两万次的水平验证集EPE也一直在原地打转。学习率、优化器、数据增强都排查过了没问题。最后定位到warping层的padding_mode。Paddle的grid_sample默认padding_modezeros也就是采样坐标超出图像边界时返回0。但在光流任务里粗层的位移经常会把采样点推出边界此时返回0意味着网络在边界附近看到的是一片黑这个虚假信号会被后续的cost volume当成真实匹配来学习。论文里的做法是用边界值填充相当于复制边缘像素也就是padding_modeborder。改这一个参数之后验证集EPE在前两万次迭代内的下降曲线明显更陡。同一时期我还发现align_corners的设置必须和损失函数里的上采样保持一致前面2.2节已经提过这里再强调一次整个流程里所有的interpolate和grid_sample要么全用True要么全用False混用会引入最多半个像素的系统偏移在光流这个以亚像素精度为目标的领域里半个像素的偏移足以让指标差出一大截。5.3 混合精度下的精度塌陷与黑名单配置为了提速我试过开自动混合精度。开启之后吞吐量确实提高了三成多但验证集EPE出现了明显的精度损失尤其在最细层的细节区域。原因是warp操作和cost volume的相关计算里存在大量接近零的小数相加fp16的精度不足以表达累积误差就上来了。解决方案是用Paddle的custom_black_list把敏感算子排除出半精度计算。我实际加进黑名单的主要是采样和相关性计算相关的算子具体名字可以通过打印算子列表确认。开启方式大致如下import paddle scaler paddle.amp.GradScaler(init_loss_scaling2 ** 16) for batch in loader: with paddle.amp.auto_cast( enableTrue, levelO1, custom_black_list{grid_sample, elementwise_mul, reduce_sum}, ): preds model(img1, img2) loss multi_scale_loss(preds, flow) scaled scaler.scale(loss) scaled.backward() scaler.step(optimizer) scaler.update()开启混合精度之后init_loss_scaling需要根据实际loss量级调太小会频繁触发跳过太大会溢出。我的经验是先从2的16次方起步观察日志里有没有大量“skipped step”的提示有就往上调一档。6. 验证、保存与断点续训的工程细节训练跑起来之后最有价值的信息来自验证集。但光流任务的验证有好几个容易写错的地方很多人训练指标和最终评测指标对不上问题往往出在这里。6.1 EPE的正确计算方式与评估时机EPEEnd-Point Error就是预测光流和真值光流之间的欧氏距离按像素平均。看起来是一行代码的事但实际计算里有三个决策点。第一哪些像素参与平均。真实数据集里存在无效区域比如遮挡、超出边界、传感器未覆盖的地方。这些像素不参与平均分母用的是有效像素数而不是总像素数。如果不做这个过滤EPE会被大量无效区域稀释看起来很低但实际没意义。第二在什么分辨率上算。论文的惯例是把预测上采样回原始图像分辨率再比较而不是在下采样后的尺寸上算。因为光流的评估是逐像素的不同分辨率下的数值不可直接比较。第三什么时候评估。训练用的裁剪尺寸和评估用的全图尺寸不一致时网络看到的感受野上下文不同性能会有偏差。我的做法是训练时用裁剪评估时也用固定裁剪但额外做一次全图评估作为参考两个数字都记录下来避免被单一指标误导。def compute_epe(pred, gt, maskNone): # pred, gt: [B, 2, H, W] epe paddle.sqrt(((pred - gt) ** 2).sum(axis1)) # [B, H, W] if mask is not None: epe epe * mask return epe.sum() / (mask.sum() 1e-6) return epe.mean()验证的频率不需要每个iteration都做我一般每五千到一万次迭代跑一次验证因为验证集全量推理本身要花几分钟。6.2 checkpoint该存什么续训该恢复什么一个完整的checkpoint至少要包含四样东西模型参数、优化器状态、学习率调度器状态、当前迭代数和最佳指标。只存模型参数是最常见的偷懒做法后果是续训之后优化器的动量估计全部丢失训练曲线会出现明显的断层。state { model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), iter: global_step, best_epe: best_epe, } paddle.save(state, fckpt_{global_step:07d}.pdparams)保存策略上我一般做两套并存每五千次迭代覆盖保存一个latest另外在验证集EPE创新低时保存一个best。这样即使某次训练中途崩溃最多损失五千次迭代而最终交付模型直接用best。续训时有一点需要注意如果恢复之后的scheduler状态和实际迭代数对不上学习率会从一个错误的位置继续后面的整个调度都错位了。恢复之后我通常会打印一下当前学习率人工确认和崩溃前的日志一致再继续跑。还有一个容易被忽略的点是最佳指标的恢复。如果best_epe没有被正确恢复续训之后可能会把之前已经更优的模型覆盖掉。这个bug很隐蔽因为训练本身完全正常只是最终产出的模型不是最好的那个。最后分享一个小技巧训练全程把每个iteration的loss、每个验证点的EPE、当前学习率都用追加模式写进一个csv不要只依赖控制台输出。我在一次长达三天的训练里控制台输出被缓冲区刷掉了一部分正是靠着csv才还原出loss在第几次迭代开始出现异常波动进而定位到是数据增强的随机种子在某次重启后没有正确恢复导致同一批样本被反复采样。光流训练的调参空间很大但真正决定成败的往往是这些看起来不起眼的工程细节一个漏乘的尺度因子、一个默认的padding模式、一个没被恢复的调度器状态。我自己的经验是每次改动之后都先用小规模数据跑个一两千次迭代看loss曲线的形状对不对再上全量。曲线形状对了后面的数字才有参考价值。
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。

↑