3D骨架舞蹈动作生成：Transformer数据管线实战

简介：AISTransformer数据代码包面向从事人工智能与机器学习数据处理的学习者和开发者，围绕原始数据向模型可训练格式的转化，提供数据清洗、特征工程与预处理等实践素材，可服务于自然语言处理、图像识别等场景的数据准备环节。压缩包共257个文件，以233个CSV数据分片为核心，配合12个Python源码文件、9个编译缓存文件及XLS、TXT、PKL等辅助文件，整体约189KB，便于查看数据组织方式并理解脚本与数据的对应关系。已有799人学习下载，具备一定参考热度。包内可重点关注Python模块与数据文件的搭配，了解预处理流程、函数与类的组织方式，以及配置、示例和测试思路，适合作为数据转换工具源码阅读与二次开发的对照材料。

1. 把 3D 骨架喂进 Transformer，卡住的从来不是模型

上个月接手一个舞蹈动作生成的活儿，数据是现成的 3D 骨架关键点加音乐节拍标注，需求也简单：给一段音乐，生成一段和节拍对得上的动作序列。我一开始的判断是模型侧最费劲，结果两周过去，写 Transformer 那部分只花了半天，剩下全耗在数据上——关键点坐标系不统一、序列长度参差不齐、节拍和帧号对不齐、滑窗切出来的片段跨了段落边界。这份 AISTransformer 数据代码解决的正是这一段：它把原始关键点、音乐特征、切片索引和 Transformer 的输入格式串成了一条插件式的管线，数据侧换数据集不用动训练脚本。适合两类人：手里有动作捕捉或视频估出来的骨架数据、想训序列模型的；以及模型已经跑通、想搞清楚工业级数据管线该怎么搭的。

2. 目录与数据流：一份舞蹈样本从 npz 文件走到 batch 的六步链路

2.1 顶层目录的分工边界

拿到包先别急着装依赖，把目录扫一遍能省掉后面很多返工。这套代码的目录划分基本遵循「原始数据只读、中间产物可重建、配置与代码分离」三条原则，动手改之前先知道哪些目录是可以整块删掉重来的。

data/processed/和data/raw/分开是血泪经验。早期我把切片写回原始目录，结果一次错误的滑窗参数把源文件覆盖了一半，只能重新从视频里估关键点。后来强制规定：raw目录挂只读，预处理产物一律落到processed，脚本里做一次路径前缀校验，写错目录直接抛异常。

2.2 一条样本的六步链路

从磁盘到模型输入，中间要过六个环节，每一步都有明确的输入输出形状约定，出问题时按这个顺序逐段打印 shape 就能定位。

第一步，__getitem__按索引从processed目录读一个 npz，里面是(T, 49, 3)的关节点坐标和(T, 35)的音乐特征。第二步，做根关节平移，把所有关节坐标减掉根关节坐标，消除全局位移。第三步，按训练/推理模式决定是随机裁一段固定长度，还是整段返回。第四步，把关节维展平成(T, 147)，和音乐特征在通道维拼接。第五步，collate_fn把同一 batch 里长度不同的序列用零补齐到 batch 内最大值，同时产出key_padding_mask。第六步，DataLoader抛出 batch，形状(B, T_max, 147+35)。

提示：第三、四步的顺序不能反。先裁窗再展平，展平后的时间维信息还能用；先展平再裁窗，很容易在 reshape 时把关节维和时间维搞混，而且这种错误不会报异常，只会让 loss 卡在某个值不动。

链路里最容易出岔子的是第四步的通道拼接。拼接顺序必须是「动作特征在前、音乐特征在后」还是反过来，取决于模型里怎么写特征切分，代码里一般用配置项feature_layout控制，别在 Dataset 里写死下标。我在一个分支上改过顺序但忘了同步改模型的切片索引，训练跑了六个小时才发现模型一直在拿音乐特征当动作标签。

2.3 插件注册表：新增数据集不碰训练脚本

关键词里带「插件」，这套代码确实做了插件式设计。核心是一个极简的注册表，数据集、损失函数、评测器三类对象都通过装饰器注册，训练脚本只按字符串名字去取，不 import 具体实现。好处是换数据集时只在plugins/dataset/下加一个文件，train.py一个字不用动。

# plugins/registry.py from typing import Callable, Dict, Any _REGISTRY: Dict[str, Dict[str, Callable]] = { "dataset": {}, # 数据集构建函数 "loss": {}, # 损失函数 "metric": {}, # 评测指标 } def register(kind: str): """带类型参数的装饰器工厂，kind 只能是 dataset / loss / metric""" if kind not in _REGISTRY: raise ValueError(f"unknown registry kind: {kind}") def wrapper(fn: Callable) -> Callable: name = fn.__name__ if name in _REGISTRY[kind]: # 同名覆盖是常见事故，这里直接报错而不是静默替换 raise KeyError(f"{kind} '{name}' already registered") _REGISTRY[kind][name] = fn return fn return wrapper def build(kind: str, name: str, **kwargs) -> Any: """按名字取实现并实例化，kwargs 来自 yaml 配置""" if name not in _REGISTRY[kind]: available = ", ".join(_REGISTRY[kind].keys()) raise KeyError(f"{kind} '{name}' not found, available: {available}") return _REGISTRY[kind][name](**kwargs)

register用工厂模式是为了支持三类对象共用一个注册表，避免写三份几乎一样的代码。build里做两层校验：名字不存在时把候选列表打出来，这在配置写错一个字母的时候特别省事；同名重复注册直接抛错而不是覆盖，因为我们确实遇到过两个分支各自注册了dance_dataset，合并时后加载的静默覆盖了前一个，训练集悄悄变成了验证集的子集。

配置侧对应写法是：

# configs/base.yaml data: name: dance_dataset # 对应 @register("dataset") 的函数名 seq_len: 150 stride: 30 train: loss: dance_loss_v2 batch_size: 32

name字段就是插件的入口点，改一行配置就换了一套数据实现，这个设计在同时维护三四条实验线的时候价值最大。

3. 数据插件拆解：滑窗采样、根关节归一化与变长补齐三件套

3.1 离线预处理脚本：把重活挪到训练之前

预处理干的事很朴素：读原始关键点，做坐标统一，按滑窗切片，落盘成 npz。之所以不在 Dataset 里实时做，是因为归一化统计量需要全量扫描一遍，实时做要么每次训练都扫一次（几十分钟），要么缓存（等价于预处理）。

# scripts/prepare_data.py import json import numpy as np from pathlib import Path JOINTS = 49 # 骨架关节数，必须和模型的输入维度对齐 SEQ_LEN = 150 # 单个切片帧数，30fps 下约 5 秒 STRIDE = 30 # 滑窗步长，约为 1 秒 EPS = 1e-6 def slice_sequence(keypoints: np.ndarray, music: np.ndarray, beats: np.ndarray, seq_len: int, stride: int): """对单条长序列做滑窗，返回片段列表，每项含起止帧号""" total = keypoints.shape[0] chunks = [] for start in range(0, total - seq_len + 1, stride): end = start + seq_len # 跳过跨越静音段（无节拍）的片段，避免生成模型学到"发呆"动作 if beats[start:end].sum() < 2: continue chunks.append({ "start": start, "end": end, "pose": keypoints[start:end].astype(np.float32), "music": music[start:end].astype(np.float32), }) return chunks def normalize_pose(pose: np.ndarray, root_idx: int = 0) -> np.ndarray: """根关节平移 + 全局尺度归一化，pose 形状 (T, J, 3)""" centered = pose - pose[:, root_idx:root_idx + 1, :] # 用整段的标准差做尺度，避免个体身高差异主导 loss scale = centered.reshape(-1, 3).std() + EPS return centered / scale def main(raw_dir: Path, out_dir: Path): out_dir.mkdir(parents=True, exist_ok=True) stats = {"count": 0, "frames": 0} for kp_file in sorted(raw_dir.glob("*.npz")): blob = np.load(kp_file) pose, music, beats = blob["pose"], blob["music"], blob["beat"] pose = normalize_pose(pose) # 音乐特征按整条序列做标准化，与动作保持同量级 music = (music - music.mean(0)) / (music.std(0) + EPS) for i, ch in enumerate(slice_sequence(pose, music, beats, SEQ_LEN, STRIDE)): np.savez_compressed( out_dir / f"{kp_file.stem}_{i:04d}.npz", pose=ch["pose"], music=ch["music"], meta=np.array([ch["start"], ch["end"]]) ) stats["count"] += 1 stats["frames"] += SEQ_LEN (out_dir / "stats.json").write_text(json.dumps(stats, indent=2))

三个参数需要说明。SEQ_LEN=150不是随便定的，它要同时满足两个约束：Transformer 的显存占用随序列长度平方增长，150 帧在单卡 24G 上 batch 32 刚好不炸；同时 5 秒的长度能覆盖一个完整的舞蹈乐句，太短的话模型学不到动作之间的过渡。STRIDE=30决定了训练样本的重叠率，步长等于帧率时重叠很少，样本量大但相邻片段差异大，收敛慢；步长太小则样本高度重复，容易过拟合到某几段。beats[start:end].sum() < 2这个过滤条件是我们后来加的，原始数据里有些片段整段是前奏或间奏，动作幅度接近零，留着会让模型在推理时倾向于生成「站着不动」的安全解。

归一化用整段的标准差而不是逐帧，是为了保留动作幅度这个信息。逐帧归一化会把大跳和手势抹成一样的幅度，生成出来的动作看着「都对但都不像」。根关节平移单独处理的原因在后面避坑章会讲，这里先记住：局部姿态和全局轨迹必须走两条通路。

3.2 Dataset 的getitem到底干了几件事

离线切好片之后，Dataset 就只剩读文件和轻量变换，但有几个细节决定训练稳不稳。

# datasets/dance_dataset.py import numpy as np import torch from torch.utils.data import Dataset from plugins.registry import register @register("dataset") class DanceDataset(Dataset): def __init__(self, index_file: str, seq_len: int = 150, training: bool = True, augment: bool = True): # 只存路径，不在这里打开任何文件句柄 self.items = [line.strip() for line in open(index_file) if line.strip()] self.seq_len = seq_len self.training = training self.augment = augment def __len__(self): return len(self.items) def __getitem__(self, idx: int): blob = np.load(self.items[idx]) # 延迟打开，worker 安全 pose = blob["pose"] # (T, J, 3) music = blob["music"] # (T, 35) if self.training and self.augment: # 时间维随机左右翻转，对应镜像动作 if np.random.rand() < 0.5: pose = pose.copy() pose[..., 0] *= -1 # 时间轴轻微抖动，提升对节拍误差的鲁棒性 shift = np.random.randint(-3, 4) if shift != 0: pose = np.roll(pose, shift, axis=0) music = np.roll(music, shift, axis=0) T, J, _ = pose.shape pose_flat = pose.reshape(T, J * 3) # (T, 147) feat = np.concatenate([pose_flat, music], axis=-1) return { "feat": torch.from_numpy(feat).float(), "length": T, }

__init__里只读索引清单、不打开数据文件，这是为多进程 DataLoader 让路——父进程打开的文件句柄 fork 到子进程后行为不可靠，运气好是读到脏数据，运气差直接报错。augment里的镜像增强要注意只翻转 x 轴坐标，不能顺手把音乐特征也翻，音乐没有镜像这个概念。时间轴np.roll的位移量控制在 ±3 帧，大约 0.1 秒，模拟的是标注时的节拍对齐误差；位移太大等于强行破坏动作和音乐的对应关系，模型会学成忽略音乐条件。

返回里带length是为了下游的 mask 和评测。别指望从feat.shape反推真实长度——补齐之后所有样本都是T_max。

3.3 collate_fn：变长序列拼 batch 的正确姿势

补齐逻辑放在 collate 而不是 Dataset，是因为只有 collate 能看到整个 batch，能做到「按 batch 内最大值补齐」而不是按全局最大值。全局补齐会让每个 batch 都为最长的那个样本付出显存。

# datasets/collate.py import torch def pad_collate(batch): lengths = torch.tensor([b["length"] for b in batch], dtype=torch.long) max_len = int(lengths.max()) dim = batch[0]["feat"].shape[-1] feats = torch.zeros(len(batch), max_len, dim) # 零填充 mask = torch.ones(len(batch), max_len, dtype=torch.bool) # True 表示 padding for i, b in enumerate(batch): L = b["length"] feats[i, :L] = b["feat"] mask[i, :L] = False return {"feat": feats, "mask": mask, "lengths": lengths}

mask的语义要提前定死：这里True表示「这是补齐出来的位置，不算数」，和 PyTorch 的key_padding_mask语义一致，直接可以传给nn.TransformerEncoder。曾经有一版反过来定义，结果注意力全屏蔽在了真实帧上，训练 loss 照样下降（因为模型学会了从零填充里提特征），只是生成的动作全是噪声。这种 bug 不看生成结果根本发现不了。

零填充配合 mask 是安全的，但要注意如果模型里有 LayerNorm 而没做 mask，补齐位置会污染统计量。稳妥做法是在进模型之前就把 padding 位置的输出丢掉再算 loss，而不是靠 LayerNorm 自己去忽略。

3.4 缓存策略与 num_workers 的取舍

num_workers设多少合适，取决于单样本读取耗时。这批 npz 单个大约 200KB，解压加反序列化不到 1ms，瓶颈在 GPU 不在 CPU，所以num_workers=4就够了。设到 16 反而因为进程创建和数据搬运的开销掉吞吐。

缓存分两层。第一层是操作系统页缓存，反复读同一批文件第二次就快了，前提是训练集能塞进内存。第二层是代码层缓存，我们试过在 Dataset 里做个 LRU 放最近读过的样本，结果因为每个 worker 各存一份、总内存是 worker 数的倍数，num_workers=8时直接吃掉十几 G。后来把缓存去掉了，改成把data/processed整个目录打成一个大文件，按偏移量读取，随机访问快很多，也不用担心句柄问题。

4. 模型与训练：编码器结构、位置编码和那几个一改就崩的超参

4.1 编码器骨架：输入投影、位置编码、层堆叠

模型侧的结构不复杂，标准的 decoder-only Transformer 换成条件生成的形式。真正决定效果的往往不是层数而是输入投影和位置编码这两处。

# models/transformer.py import math import torch import torch.nn as nn class SinusoidalPE(nn.Module): def __init__(self, d_model: int, max_len: int = 4096): super().__init__() pe = torch.zeros(max_len, d_model) pos = torch.arange(0, max_len).unsqueeze(1).float() div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(pos * div) pe[:, 1::2] = torch.cos(pos * div) self.register_buffer("pe", pe) # 不参与训练，跟模型一起搬设备 def forward(self, x: torch.Tensor) -> torch.Tensor: return x + self.pe[: x.size(1)] class DanceTransformer(nn.Module): def __init__(self, pose_dim: int = 147, music_dim: int = 35, d_model: int = 512, nhead: int = 8, num_layers: int = 6, dropout: float = 0.1): super().__init__() self.in_proj = nn.Linear(pose_dim + music_dim, d_model) self.pe = SinusoidalPE(d_model) enc_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4, dropout=dropout, batch_first=True, norm_first=True) self.encoder = nn.TransformerEncoder(enc_layer, num_layers=num_layers) self.norm = nn.LayerNorm(d_model) self.out_proj = nn.Linear(d_model, pose_dim) def forward(self, feat, key_padding_mask=None): h = self.in_proj(feat) # (B, T, D) h = self.pe(h) # 加正弦位置编码 h = self.encoder(h, src_key_padding_mask=key_padding_mask) h = self.norm(h) # 出栈后再归一化，稳定输出尺度 return self.out_proj(h) # (B, T, 147)

in_proj把 147 维姿态和 35 维音乐拼成 182 维再投到d_model，这个投影层看着不起眼，但它是动作和音乐两种模态第一次融合的地方。我们试过两个模态各走一个线性层再相加，效果明显差于先拼接再投影——拼接保留了两种特征在同一帧上的共现关系，而加和会把这层关系搅乱。

norm_first=True是后来的改动。原始参数是 post-LN，训到十几个 epoch 之后 loss 偶尔会突然跳起来，换成 pre-LN 之后稳定不少，代价是收敛稍慢一点。dim_feedforward保持 4 倍是常规配置，改小到 2 倍能省显存但表达力明显下降，改大到 8 倍在 6 层的情况下收益很小。

位置编码用正弦而不是可学习，主要是为了支持推理时长序列外推。可学习的位置编码在超过训练长度之后没有定义，而舞蹈生成经常需要生成比训练片段更长的序列。

4.2 训练循环：三个损失项各管一段

损失不是单一 MSE。纯 MSE 训出来的动作逐帧都「平均正确」，但播放起来抖得厉害，因为相邻帧之间没有约束。

# train.py 核心片段 import torch def dance_loss(pred, gt, mask, w_rec=1.0, w_vel=0.5, w_acc=0.1): valid = ~mask # True 表示真实帧 # 1) 逐帧重建误差 rec = ((pred - gt) ** 2).mean(-1)[valid].mean() # 2) 速度一致性：相邻帧位移的一阶差分 v_pred = pred[:, 1:] - pred[:, :-1] v_gt = gt[:, 1:] - gt[:, :-1] vel_mask = valid[:, 1:] & valid[:, :-1] vel = ((v_pred - v_gt) ** 2).mean(-1)[vel_mask].mean() # 3) 加速度平滑项，压制高频抖动 a_pred = v_pred[:, 1:] - v_pred[:, :-1] acc_mask = vel_mask[:, 1:] & vel_mask[:, :-1] acc = (a_pred ** 2).mean(-1)[acc_mask].mean() return w_rec * rec + w_vel * vel + w_acc * acc, { "rec": rec.item(), "vel": vel.item(), "acc": acc.item()}

三个权重不是拍脑袋定的：w_rec=1.0是基准，w_vel=0.5让模型关注运动趋势，w_acc=0.1只做轻微平滑。第二项调大到 2.0 以上会出现动作「拖影」，因为模型倾向于用匀速直线运动去拟合所有轨迹；第三项调大到 0.5 以上动作会变得僵硬，像机器人。这两个数值我改过至少五轮，最后定在现在的组合。

mask 的处理是这段代码里最容易写错的地方。vel_mask是两个相邻有效位的与，最后一个 padding 位置和第一个 padding 位置永远不参与计算，否则补齐的零会让速度差分变成真实帧坐标本身，产生巨大的虚假梯度。

训练循环其它部分就常规了，AdamW + cosine 衰减，warmup 2000 步，混合精度。混合精度下要留意 loss 里如果出现-inf（比如 mask 全被屏蔽时mean()返回 nan），会被 GradScaler 直接吞掉表现为 loss 不动，加一行torch.nan_to_num或者断言能早点暴露。

4.3 评测脚本：别只看 FID

评测用两类指标。一类是分布距离，把生成动作和真实动作都送进一个预训练的动作特征提取器，算特征分布之间的 Frechet 距离，数值越低越接近真实分布。另一类是多样性，同一段音乐采多个样本，算两两之间的平均距离，太低说明模型只会生成一种动作，太高说明生成的东西和音乐没关系。

关键参数是采样数。生成 50 个样本算出来的分布距离方差很大，同一模型两次评测能差出 10%，所以固定用 200 个样本、固定随机种子。评测时还要把生成结果和真实数据用同一套归一化参数处理，曾经因为评测脚本忘了反归一化，FID 直接涨了三倍，白折腾了两天。

5. 避坑与排查：五个让我半夜重跑训练的现象

5.1 数据与路径相关的三个坑

现象一：训练 loss 第一个 epoch 就掉到 0.01 以下，但生成的动作像抽风。原因：滑窗切片时把start和end的索引写反了，keypoints[start:end]变成了keypoints[end:start]，返回空数组，npz 里存的是上一轮的旧数据。Dataset 读到的样本和自己的标签来自不同文件，模型学的是「任意输入映射到固定输出」。这类错误不会抛异常，因为形状是对的。 解决：预处理脚本末尾加一条断言，检查每个输出文件的meta字段里的帧号范围和pose.shape[0]一致，不一致直接中止。另外在 Dataset 里加一个__getitem__的样本自检开关，训练前抽 3 个样本把动作序列画出来看一眼。

现象二：多进程加载偶发FileNotFoundError或者读到全零。原因：在Dataset.__init__里用np.load打开了所有文件并常驻内存，父进程 fork 出 worker 后句柄失效。在num_workers=0时一切正常，改成 8 就开始随机出错，非常难复现。 解决：__init__只存路径，__getitem__里每次打开。如果确实要做全量缓存，改成用multiprocessing的共享内存或者干脆在worker_init_fn里重新打开。

现象三：生成的舞蹈原地踏步，手臂动作对但整个人不移动。原因：3.2 里做的根关节平移把全局位移信息抹掉了，模型拿到的输入里根本没有「人在往哪走」这个信号。这个问题在训练 loss 上看不出来，因为局部姿态确实拟合得很好。 解决：把根关节轨迹单独拉出来做一路回归，和局部姿态两路输出，loss 里加一项轨迹的 L1 误差，权重设 0.2 左右。或者干脆不对根关节做平移，只做尺度归一化，让模型自己学。

5.2 训练与显存相关的两个坑

现象四：batch_size设成 8 还是 OOM，但显存监控显示占用忽高忽低。原因：pad_collate按 batch 内最大长度补齐，数据集里有少数几个超长片段（预处理时没按seq_len截断的残留），一旦被采进 batch，显存峰值就翻好几倍。shuffle=True让这个问题以随机频率出现，看起来像显存泄漏。 解决：预处理阶段加长度校验，超过seq_len * 1.2的直接丢弃；同时在 collate 里加一个硬上限裁剪，超过上限的序列直接截断并同步修正 mask 和 length。

现象五：混合精度开启后 loss 变成 nan，关掉就正常。原因：nn.TransformerEncoder在 fp16 下注意力分数容易溢出，尤其是位置编码加得很满、输入值域偏大的时候。另外dance_loss里的除法在 mask 全 False 时会产生 0/0。 解决：编码器整体用torch.autocast包住，但LayerNorm和 loss 计算强制走 fp32；输入先做一次clamp(-10, 10)防止离群值；mask 求和前加max(1, ...)兜底。

6. 二次开发：换数据集、换编码器、加一个自定义损失插件

这套代码最容易改的地方和最容易改坏的地方是同一个——插件层。换数据集只需要三步：在plugins/dataset/下新建文件，写一个带@register("dataset")的类，保证__getitem__返回{"feat": (T, D), "length": int}；然后改 yaml 里的data.name和data.in_dim；最后跑一遍预处理脚本产出对应的 index 文件。整个过程不碰train.py。

换编码器同理，但有个隐含契约要满足：输入(B, T, D_in)，输出(B, T, 147)，中间自己处理key_padding_mask。如果想换成卷积或者 Mamba 类的结构，需要注意掩码语义是否一致，别直接把 Transformer 的 mask 传进去。

加一个损失插件的写法：

# plugins/loss/foot_contact.py import torch from plugins.registry import register @register("loss") def foot_contact_loss(foot_idx=0, thr=0.02, weight=0.05): """脚部接触约束：脚落地时应贴近地面，防止滑步""" def _loss(pred, gt, mask): valid = ~mask foot_y = pred[..., foot_idx * 3 + 1] # 取根关节的 y 分量 # 只惩罚"应该落地但抬起来"的帧 gt_foot_y = gt[..., foot_idx * 3 + 1] contact = (gt_foot_y < thr) & valid if contact.sum() == 0: return pred.new_zeros(()) return weight * (foot_y[contact].abs()).mean() return _loss

注意这个损失返回的是闭包，build("loss", "foot_contact_loss", weight=0.1)拿到的才是可调用的损失函数，配置里的权重和函数签名要能对上，否则build会在初始化阶段就报 TypeError，比训练到一半才炸好得多。

验证改造是否成功，我的习惯是固定跑三个小实验：把seq_len减到 30 看能不能过拟合 8 个样本（验证数据管线没断），把num_layers设成 1 看 loss 是否还能正常下降（验证梯度通路），把评测集换成训练集的一个子集看指标是否接近完美（验证评测脚本没写反）。这三个跑通，再上全量训练。

注意：改任何和数据形状相关的参数之后，先把data/processed/整个删掉重跑预处理。增量更新在这里是陷阱，旧切片和新参数混在一个目录里，训练时不会报错，只会在评测时表现为指标莫名其妙地差。

从那以后我每次接手一个带离线预处理的项目，第一件事就是把预处理脚本从头到尾跑一遍并且断点检查中间产物的形状，第二步才是装依赖跑训练——因为模型层面的错误通常有异常堆栈，数据层面的错误往往只有一张看起来还行但说不出哪不对的生成结果。希望帮到你。

本文还有配套的精品资源，点击获取