简介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/raw/原始关键点、音乐特征、节拍标注只读几乎不改data/processed/预处理后的 npz 切片、归一化统计量可重建随便删data/splits/训练/验证/测试的片段 ID 清单换数据集时改datasets/Dataset 类、collate、预处理脚本常改models/Transformer 编码器、位置编码、输出头常改plugins/数据集/损失/评测的注册表实现扩展时加configs/yaml 配置按实验分文件每次实验都改scripts/预处理、训练、评测、可视化的入口偶尔改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, 14735)。提示第三、四步的顺序不能反。先裁窗再展平展平后的时间维信息还能用先展平再裁窗很容易在 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(funknown 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: 32name字段就是插件的入口点改一行配置就换了一套数据实现这个设计在同时维护三四条实验线的时候价值最大。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(parentsTrue, exist_okTrue) 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, posech[pose], musicch[music], metanp.array([ch[start], ch[end]]) ) stats[count] 1 stats[frames] SEQ_LEN (out_dir / stats.json).write_text(json.dumps(stats, indent2))三个参数需要说明。SEQ_LEN150不是随便定的它要同时满足两个约束Transformer 的显存占用随序列长度平方增长150 帧在单卡 24G 上 batch 32 刚好不炸同时 5 秒的长度能覆盖一个完整的舞蹈乐句太短的话模型学不到动作之间的过渡。STRIDE30决定了训练样本的重叠率步长等于帧率时重叠很少样本量大但相邻片段差异大收敛慢步长太小则样本高度重复容易过拟合到某几段。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, axis0) music np.roll(music, shift, axis0) 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], dtypetorch.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, dtypetorch.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_workers4就够了。设到 16 反而因为进程创建和数据搬运的开销掉吞吐。缓存分两层。第一层是操作系统页缓存反复读同一批文件第二次就快了前提是训练集能塞进内存。第二层是代码层缓存我们试过在 Dataset 里做个 LRU 放最近读过的样本结果因为每个 worker 各存一份、总内存是 worker 数的倍数num_workers8时直接吃掉十几 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_modeld_model, nheadnhead, dim_feedforwardd_model * 4, dropoutdropout, batch_firstTrue, norm_firstTrue) self.encoder nn.TransformerEncoder(enc_layer, num_layersnum_layers) self.norm nn.LayerNorm(d_model) self.out_proj nn.Linear(d_model, pose_dim) def forward(self, feat, key_padding_maskNone): h self.in_proj(feat) # (B, T, D) h self.pe(h) # 加正弦位置编码 h self.encoder(h, src_key_padding_maskkey_padding_mask) h self.norm(h) # 出栈后再归一化稳定输出尺度 return self.out_proj(h) # (B, T, 147)in_proj把 147 维姿态和 35 维音乐拼成 182 维再投到d_model这个投影层看着不起眼但它是动作和音乐两种模态第一次融合的地方。我们试过两个模态各走一个线性层再相加效果明显差于先拼接再投影——拼接保留了两种特征在同一帧上的共现关系而加和会把这层关系搅乱。norm_firstTrue是后来的改动。原始参数是 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_rec1.0, w_vel0.5, w_acc0.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_rec1.0是基准w_vel0.5让模型关注运动趋势w_acc0.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_workers0时一切正常改成 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显存峰值就翻好几倍。shuffleTrue让这个问题以随机频率出现看起来像显存泄漏。 解决预处理阶段加长度校验超过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 传进去。改造目标改动位置需要同步改的配置换骨架关节点数prepare_data.py的JOINTS、模型pose_dimdata.in_dim、输出维度换音乐特征prepare_data.py特征拼接段、music_dimdata.music_dim加自定义损失plugins/loss/新文件 train.py权重参数train.loss、train.w_*换位置编码models/transformer.py的SinusoidalPE无需改配置但需重训加一个损失插件的写法# plugins/loss/foot_contact.py import torch from plugins.registry import register register(loss) def foot_contact_loss(foot_idx0, thr0.02, weight0.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, weight0.1)拿到的才是可调用的损失函数配置里的权重和函数签名要能对上否则build会在初始化阶段就报 TypeError比训练到一半才炸好得多。验证改造是否成功我的习惯是固定跑三个小实验把seq_len减到 30 看能不能过拟合 8 个样本验证数据管线没断把num_layers设成 1 看 loss 是否还能正常下降验证梯度通路把评测集换成训练集的一个子集看指标是否接近完美验证评测脚本没写反。这三个跑通再上全量训练。注意改任何和数据形状相关的参数之后先把data/processed/整个删掉重跑预处理。增量更新在这里是陷阱旧切片和新参数混在一个目录里训练时不会报错只会在评测时表现为指标莫名其妙地差。从那以后我每次接手一个带离线预处理的项目第一件事就是把预处理脚本从头到尾跑一遍并且断点检查中间产物的形状第二步才是装依赖跑训练——因为模型层面的错误通常有异常堆栈数据层面的错误往往只有一张看起来还行但说不出哪不对的生成结果。希望帮到你。本文还有配套的精品资源点击获取