
简介这是一份基于PyTorch框架、面向时间序列轨迹预测的Transformer项目源码包适合具备一定深度学习基础、想通过完整案例上手机器学习预测的开发者与研究生。项目自带可直接运行的主程序、训练脚本和多种格式的轨迹数据覆盖数据清洗、特征构造、模型定义、训练评估与结果可视化全流程。压缩包共58个文件、大小约312MB主要由12个py源码文件、5个ipynb交互式笔记本、7个tar模型权重、2个xlsx和2个csv数据集及相关配置组成便于按需查看和断点调试。目前已有918人浏览学习是研究自注意力机制在时间序列中应用的实用素材。运行项目可观察Transformer如何建模历史轨迹与未来位置的关系掌握基于PyTorch的Transformer搭建、超参数调整和模型评估方法为后续交通、机器人等场景的轨迹预测提供可复现的参考。 Transformer做轨迹预测这件事我在不同项目里前前后后折腾了大半年从最初照搬NLP里的标准实现到后面针对轨迹数据的特点一点点改结构、调掩码、换位置编码踩过的坑比写过的代码还多。这次这个Pytorch实现我把数据集生成、模型搭建、训练推理的完整代码都整理到了一起解压之后按顺序跑就能看到预测效果。如果你正在做车辆轨迹预测、行人路径规划或者想用Transformer处理序列回归问题这份代码应该能帮你省下不少排查bug的时间。1. 为什么轨迹预测要用Transformer而不是LSTM或CNN先聊一个最基础的问题轨迹预测本质上是什么它是在给定历史轨迹的条件下预测未来的位置序列。传统做法里LSTM是按时间步逐个处理CNN是用卷积核在时间维度上滑动它们在处理短序列时表现尚可但一旦轨迹变长、agent数量变多问题就出来了。Transformer和它们最大的区别在于注意力机制。自注意力让序列里任意两个时间步之间都能直接建立联系这意味着模型可以同时捕捉“5帧之前的位置变化趋势”和“当前帧的瞬时状态”之间的关系不会像RNN那样信息逐层衰减。另一个更关键的优势是轨迹预测往往涉及多个目标之间的交互。比如十字路口的两辆车它们的轨迹是互相影响的。Transformer的编码器天然支持把多个agent的轨迹拼成一个batch一起编码注意力权重会把“谁在影响谁”隐式地学出来。这一点单靠LSTM需要额外设计复杂的交互模块才能做到。当然这不代表Transformer在所有场景下都吊打RNN。我在小规模数据集上对比过如果只有单目标、轨迹长度很短比如10帧以内LSTM的效果和Transformer差距不大但Transformer的收敛速度更快而且不需要手动设计时序特征。如果你正在处理多目标交互场景或者数据量足够大Transformer几乎是最省心的选择。2. 数据集的构造逻辑与代码细节这份代码里附带的数据集是合成的但它的格式和真实业务场景完全一致每一行代表一个时间帧包含timestamp, agent_id, x, y四个字段。保存在data/raw/目录下名为trajectory_data.csv。之所以用合成数据是为了让你在没有真实业务数据的情况下也能快速跑通全流程替换成自己的数据只需要保持字段名不变就行。合成数据的生成逻辑其实很能说明问题它模拟了三种运动模式——直线匀速、匀速转弯、以及两辆车互相靠近的交互场景。第三类数据是重点因为Transformer在交互预测上的优势主要靠这部分体现。生成之后数据会经过一个预处理脚本切成(history_len, future_len, 2)形状的样本其中history_len默认是8帧future_len是12帧2代表x和y坐标。这里有一个值得注意的细节原始坐标必须做归一化。我用的方法是对每个轨迹样本单独计算均值和标准差然后做标准化。这一步不做的话模型在训练时容易因为坐标尺度过大比如像素坐标上千导致梯度爆炸或者在预测时输出一些离谱的大数值。归一化时要注意测试阶段也使用训练阶段保存的均值和标准差不能重新计算否则效果会崩。数据加载部分我写了一个继承torch.utils.data.Dataset的类里面做了两件容易被忽略的事一是把所有轨迹样本补齐到相同长度用mask标记真实数据和pad部分二是把相邻agent的轨迹关联信息生成出来作为辅助输入。第一点尤其重要因为torch的DataLoader要求batch内所有样本形状一致很多人第一次跑Transformer序列任务就卡在这里。3. 模型结构从嵌入层到解码器的实现这套Transformer轨迹预测模型不是简单套用nn.Transformer而是针对轨迹序列做了定制主要包含四个模块输入嵌入层、位置编码、编码器、解码器。我建议你直接看models/transformer_trajectory.py这个文件下面讲几个核心实现细节。3.1 输入嵌入层的维度设计原始输入是(batch_size, history_len, 2)的坐标序列首先通过一个nn.Linear(2, d_model)把坐标映射到高维空间。d_model我默认设成128这个值对大多数轨迹预测任务都够用。除了坐标本身我把速度信息也拼进了嵌入向量先对原始序列做一阶差分得到速度然后也过一个线性层最后和位置嵌入相加。这么做的好处是模型不需要自己从坐标里隐式学速度降低了训练难度。实测下来加了速度分支后预测平滑度有明显提升。3.2 位置编码为什么不能直接照搬NLP版本这是我在实践中踩过最大的一次坑。NLP里的标准位置编码是sinusoidal的因为token的绝对位置对语言理解很重要。但轨迹数据完全不同它本质上是一条平滑的曲线相邻帧之间的位移很小而且轨迹的起点位置是随机的在不同样本里可能完全不同。如果沿用NLP的位置编码模型会倾向于记住“第5帧的位置应该在哪”而不是“第5帧相对于第1帧的位移是多少”。我最后的做法是位置编码改为归一化的相对时间编码。具体来说对每个样本把时间步索引除以该样本的总步数得到一个0到1之间的小数然后过一个可学习的nn.Embedding层与输入嵌入相加。这样模型学到的是“当前处于整条轨迹的什么进度”而不是某个绝对的时间点。这个改动让测试集上的ADE平均位移误差Average Displacement Error下降了约18%效果非常明显。3.3 编码器与解码器的掩码设置编码器标准实现使用的是nn.TransformerEncoder和nn.TransformerEncoderLayer。但注意默认的编码器是双向注意力即每个位置都能看到整条历史序列的所有位置这对轨迹编码是合理的——我们就是想model整段历史轨迹的全局特征。解码器用的自回归模式推理时要一个点一个点地预测未来轨迹。这里必须用nn.TransformerDecoder配合generate_square_subsequent_mask来防止未来信息泄漏。很多人直接拿编码器当解码器用预测质量会差不少因为训练和推理行为不一致。另外一个关键点是src_key_padding_mask。我们在数据预处理时做了paddingpadding部分必须在注意力里被屏蔽掉否则pad位置会参与计算污染真实轨迹的特征。代码里create_padding_mask函数可以自动生成这个掩码训练和推理时都要记得传入。3.4 解码策略还是递归预测更稳轨迹预测的解码方式有两种一次性输出全部未来轨迹并行解码还是一步步递归预测自回归解码。我两种都试过。并行解码快训练也简单直接设定输出长度即可但实际效果不稳定经常出现轨迹抖动。最终我采用了自回归方式把历史轨迹的最后一个位置作为第一个预测输入然后拿预测结果作为下一步的输入循环12次得到完整未来轨迹。这个方式和真实物理世界更贴合——目标每一帧的位置变化是逐步推进的递归预测更符合动力学的连续演化规律。训练时我用了teacher forcing即输入真实的前一帧位置推理时则完全用自己的预测结果作为下一步输入。4. 训练流程、损失函数与超参调优训练脚本在main.py里核心部分不复杂但有几个参数直接影响效果我把自己的调参经验写在下面。4.1 训练循环怎么写标准训练循环逻辑数据加载、前向传播、计算损失、反向传播、梯度裁剪、优化器step。我额外加了梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。不加这一行训练过程中loss曲线会出现尖刺尤其在序列较长时。循环里每50步打印一次当前epoch、step、loss。模型保存用的是torch.save(model.state_dict(), path)检查点保存在checkpoints/best_model.pth。我会在验证集上计算ADE只有当当前模型在验证集上的ADE优于历史最优时才覆盖保存这样能避免训练后期过拟合导致模型反而变差。4.2 损失函数不仅是MSE轨迹预测最常用的评估指标是ADE平均位移误差和FDE最终位移误差Final Displacement Error。ADE计算的是整条预测轨迹每个点和真实轨迹对应点的平均距离FDE只看最后一步的偏差距离。损失函数我直接用MSE均方误差即预测点和真实点的欧氏距离平方这和ADE在本质上是正相关的。但单用MSE有个问题预测轨迹往往偏保守倾向于输出平均位置导致在转弯场景下预测轨迹“抄近路”把弯道拉直了。后来我在损失函数里加了一项平滑约束相邻两个预测点之间的距离差应当接近相邻两个真实点之间的距离差。这个约束本质上是让预测轨迹保持合理的运动学特征不至于出现相邻帧位移剧烈跳变的异常情况。加了这个平滑项后L2损失只下降了约5%但FDE下降了12%左右明显改善了对转弯场景的拟合。4.3 这份配置在测试集上的成绩我直接给出这份代码默认配置下的测试结果方便你验证自己的改动是否有效指标默认配置ADE平均位移误差0.731FDE最终位移误差1.285训练时间CPU500个epoch约18分钟训练时间GPU约4分钟默认参数如下batch_size128lr1e-3d_model128nhead4num_encoder_layers3num_decoder_layers3dropout0.1history_len8future_len12。如果想更快看到效果可以适当增大batch_size或者把num_encoder_layers减到2层。4.4 优化器与学习率调度优化器用的是AdamWweight_decay设为1e-5比Adam多一个解耦的权重衰减在小数据集上不容易过拟合。学习率调度我用的是一阶段余弦退火CosineAnnealingLRT_max设为总epoch数。前50步做warmup是因为Transformer对初始学习率比较敏感warmup能让模型在早期有个稳定的上升过程之后余弦退火慢慢降低学习率在训练后期做细粒度收敛。这套组合亲测比固定学习率稳定很多。5. 推理脚本与可视化怎样把预测结果变成直观的图模型训练好之后光看loss数字是没有感觉的我建议一定要做可视化。predict.py这个脚本就是干这件事的加载测试数据、载入训练好的模型权重、输出预测轨迹并用matplotlib画出历史轨迹、真实未来轨迹和预测轨迹三条线。运行方式python predict.py --checkpoint checkpoints/best_model.pth --data data/processed/test.csv --output results/prediction.png可视化脚本里我用不同颜色区分了三类轨迹灰色是历史轨迹绿色是真实未来轨迹红色是模型预测的未来轨迹。这段逻辑不复杂实现上就是把历史坐标、真实未来坐标、预测输出三条线依次画出来。我在实际跑推理时发现一个常见问题有些样本预测轨迹的前几步和真实轨迹比较接近但越到后面偏得越远。这大概率是训练时teacher forcing和推理时完全自回归之间的gap造成的。改进方向有两个一是训练末期逐渐降低teacher forcing的概率比如从1.0线性降到0.5让模型逐步适应推理阶段的输入分布二是在推理时引入一个小范围的高斯噪声把模型预测出的轨迹略微平滑后再作为下一步输入相当于做了一次隐式平滑能减少误差累积。另外推理脚本里我加了GPU和CPU的自动检测有CUDA就用GPU没有就自动切回CPU。所以即使你机器上没有NVIDIA显卡也一样能跑通整个流程只是速度会慢一些。6. 跑通Demo之后的经验总结与踩坑记录最后聊一些代码之外的体会这些是我在调这个模型时花时间最多的地方也希望你少走弯路。第一个坑是位置编码。我在第3.2节已经详细说过标准NLP的位置编码不能用必须改成基于归一化相对位置的编码。这个问题的排查过程让我印象很深一开始模型loss能正常下降但可视化出来的预测轨迹完全“没形”就像是模型把历史轨迹复制了一份然后强行平移。后来我打印了中间层输出发现编码器的表示里带有很强的绝对位置偏好改了位置编码之后才好转。如果你遇到类似症状先检查位置编码方案而不是盲目加深网络。第二个坑是padding掩码。batch里轨迹长度不一致padding成相同长度后如果你忘了传入src_key_padding_mask注意力会分一部分权重给那些纯零的位置导致每个样本的特征都被“稀释”。这个问题的表现形式是loss死活降不到预期水平。用上面提到的create_padding_mask生成mask并在transformer.encoder和transformer.decoder里都传入能直接改善2%~3%的测试指标。第三个坑是数据顺序。我一开始做数据预处理时对轨迹样本做了一个全局shuffle导致同一个agent的历史轨迹和未来轨迹被打散了。这听起来像是低级错误但真实项目里数据管道一复杂就会出这种问题。训练数据里编码器的输入是历史轨迹解码器的监督是未来轨迹两者必须严格按时间顺序配对。建议在Dataset类内部维护好索引不要在外部对csv行做随机重排。第四个坑是关于模型大小和数据集规模的匹配。Transformer不是越深越好。我在只有几百个样本的数据集上试过6层编码器结果直接过拟合训练集loss降到很低验证集曲线乱跳。这份代码默认3层编码器3层解码器是在小规模数据下比较均衡的选择。如果你的数据量超过几万条再去加深网络收益会明显很多。最后再分享一个小技巧如果你要在训练期间看模型到底学到了什么每训练几十个epoch就挑一个固定测试样本跑一次推理、存一张可视化图按epoch编号命名整理在一起。你会非常直观地看到预测轨迹从“乱飞”到“贴合”的过程这种反馈对调参的价值比看一堆loss曲线大得多。我在开发这套代码时一直保留这个习惯很多玄学调参问题都是靠盯图盯出来的。这套代码后续还可以往几个方向扩展比如把多个agent之间的交互关系通过一个额外的注意力层显式建模出来或者接入高清地图数据做车道约束预测。如果你之前只用LSTM做过轨迹预测我建议花一个下午跑通这份代码感受一下注意力机制在处理空间-时间序列时的真正实力。本文还有配套的精品资源点击获取