项目概览
bytedance/piano_transcription 是字节跳动提出的高分辨率钢琴转写系统（High-resolution Piano Transcription with Pedals）的官方 PyTorch 实现，能够将钢琴演奏录音自动转写为包含音符（音高、起音、止音、力度）与踏板信息的 MIDI 文件。本页为整个代码库的总览入口，介绍项目定位、总体架构、核心模块划分与端到端工作流。
目的与范围（Purpose and Scope）
本页覆盖以下内容：
- 项目的研究背景与任务定义（钢琴转写：音频 → MIDI）
- 代码库的整体架构与模块划分（
pytorch/训练推理核心、utils/数据与特征工具、predict.pyCog 推理接口） - 基于论文 [1] 的两阶段训练策略（音符模型 + 踏板模型 → 合并）与完整流水线（runme.sh）
- CRNN 声学模型的基本结构（
ConvBlock/AcousticModelCRnn8Dropout） - 训练环境、数据集（MAESTRO V2.0.0）、资源需求与基线结果
以下主题属于姊妹页面，本页仅作指引、不深入展开：
- 推理细节（
pytorch/inference.py的转写与后处理流程）——参见推理相关页面 - 训练循环、损失函数与优化器配置——参见训练与损失相关页面
- 数据打包、特征提取与数据增强——参见数据处理相关页面
- 评测指标计算（
calculate_score_for_paper.py）——参见评测相关页面
概述（Overview）
钢琴转写 是音乐信息检索（MIR）中的经典任务：输入一段钢琴演奏音频，输出符号化的音符序列（MIDI）。与通用自动语音识别不同，钢琴转写需要同时预测：
| 输出目标 | 含义 | 建模范式 |
|---|---|---|
| onset（起音） | 音符开始时刻 | 回归精确时间（高分辨率） |
| offset（止音） | 音符结束时刻 | 回归精确时间（高分辨率） |
| frame（帧激活） | 每帧是否有音发声 | 二分类 |
| velocity（力度） | 音符击键强度 | 回归 |
| pedal（踏板） | 延音/柔音踏板的按下与抬起 | 回归精确时间 |
本项目的核心创新（论文 [1]）在于通过回归精确的起音/止音时间来提升时间分辨率，而不是传统上仅做逐帧二分类再取帧中心作为音符边界。这带来 10 ms 级别的音符边界精度，并在 MAESTRO 测试集上达到 note F1 = 0.9677、pedal F1 = 0.9186 的水平（见预训练 checkpoint 命名 CRNN_note_F1=0.9677_pedal_F1=0.9186.pth）。
典型使用场景：
- 开箱即用推理：通过
piano_transcription_inferencepip 包（或本仓库pytorch/inference.py）加载预训练模型，将 mp3/wav 转为 MIDI。 - 从头训练：在 MAESTRO V2.0.0 数据集上分别训练音符系统与踏板系统，再合并权重。
- 大规模应用：官方基于本系统构建了大规模古典钢琴 MIDI 数据集 GiantMIDI-Piano（见 README "Applications" 一节）。
运行环境：Python 3.7 + PyTorch 1.4.0 开发（README 声明其他版本应可用但未完整测试）；训练需要约 29 GB 显存（batch size = 12，单卡 Tesla-V100-PCIE-32GB 训练 300k 迭代约一周）。
架构（Architecture）
整个代码库由三层组成：数据层（MAESTRO 数据集 → hdf5 打包 → 特征提取）、模型层（CRNN 声学模型 + 多任务回归头）、应用层（训练脚本、推理脚本、评测脚本、Cog 在线推理接口）。总体数据流如下：
Loading diagram...
分层设计意图：
- 数据层与模型层解耦：
utils/data_generator.py负责从 hdf5 中随机裁剪训练片段并在线计算特征，使模型层只关心张量进张量出，便于切换数据集（runme.sh 中的评测同时支持maestro与maps两个数据集）。 - 音符与踏板两套模型、一个骨干：训练阶段分别以
Regress_onset_offset_frame_velocity_CRNN（音符）和Regress_pedal_CRNN（踏板）两种model_type训练，共享同一 CRNN 骨干结构；训练完成后由combine_note_and_pedal_models.py将两份 checkpoint 合并为单一Note_pedal模型，推理时一次前向同时输出音符与踏板。 - 推理独立成包：日常使用无需本仓库训练代码，官方另发布了
piano_transcription_inferencepip 包；本仓库的predict.py则是面向 Replicate 平台的 Cog 封装。
模块组成
代码库顶层结构（基于实际文件列表）：
| 路径 | 角色 | 说明 |
|---|---|---|
pytorch/main.py | 训练入口 | 解析命令行参数，执行 train 模式，按 iteration 保存 checkpoint 与统计量 |
pytorch/models.py | 模型定义 | ConvBlock、AcousticModelCRnn8Dropout 及各类 Regress_* 顶层模型 |
pytorch/losses.py | 损失函数 | 与 loss_type 对应的多任务回归损失 |
pytorch/inference.py | 本地推理 | 加载合并后的 Note_pedal checkpoint，音频 → MIDI |
pytorch/combine_note_and_pedal_models.py | 权重合并 | 把音符 checkpoint 与踏板 checkpoint 合并 |
pytorch/evaluate.py / calculate_score_for_paper.py | 评测 | infer_prob + calculate_metrics 两段式评测 |
pytorch/pytorch_utils.py | 训练工具 | 数据搬移到设备等辅助函数 |
utils/data_generator.py | 数据生成器 | 从 hdf5 采样训练片段 |
utils/features.py | 特征 | 频谱/对数梅尔特征计算 |
utils/config.py | 路径配置 | workspace、数据集路径解析 |
utils/piano_vad.py | VAD | 钢琴静音段检测辅助 |
utils/utilities.py, utils/plot_statistics.py, utils/plot_for_paper.py | 工具/绘图 | 统计曲线与论文图表 |
predict.py | Cog 推理接口 | Replicate 平台的在线预测封装 |
runme.sh | 一键流水线 | 从数据打包到评测的完整命令序列 |
CRNN 声学模型骨架
pytorch/models.py 中的声学模型是一个卷积循环网络（CRNN）。其前端的频谱变换来自 torchlibrosa：
python
from torchlibrosa.stft import Spectrogram, LogmelFilterBank
from pytorch_utils import move_data_to_deviceSource: pytorch/models.py
基础卷积块 ConvBlock 由两个 3×3 卷积 + BatchNorm + ReLU + 平均池化构成，这是整个网络的空间特征提取单元：
python
1class ConvBlock(nn.Module):
2 def __init__(self, in_channels, out_channels, momentum):
3
4 super(ConvBlock, self).__init__()
5
6 self.conv1 = nn.Conv2d(in_channels=in_channels,
7 out_channels=out_channels,
8 kernel_size=(3, 3), stride=(1, 1),
9 padding=(1, 1), bias=False)
10
11 self.conv2 = nn.Conv2d(in_channels=out_channels,
12 out_channels=out_channels,
13 kernel_size=(3, 3), stride=(1, 1),
14 padding=(1, 1), bias=False)
15
16 self.bn1 = nn.BatchNorm2d(out_channels, momentum)
17 self.bn2 = nn.BatchNorm2d(out_channels, momentum)
18
19 self.init_weight()Source: pytorch/models.py
骨干网络 AcousticModelCRnn8Dropout 堆叠 4 个 ConvBlock（通道数 1 → 48 → 64 → 96 → 128），经全连接降维到 768 后接 2 层双向 GRU（hidden 256，双向拼接为 512），最后由线性层映射到 classes_num 个输出目标：
python
1class AcousticModelCRnn8Dropout(nn.Module):
2 def __init__(self, classes_num, midfeat, momentum):
3 super(AcousticModelCRnn8Dropout, self).__init__()
4
5 self.conv_block1 = ConvBlock(in_channels=1, out_channels=48, momentum=momentum)
6 self.conv_block2 = ConvBlock(in_channels=48, out_channels=64, momentum=momentum)
7 self.conv_block3 = ConvBlock(in_channels=64, out_channels=96, momentum=momentum)
8 self.conv_block4 = ConvBlock(in_channels=96, out_channels=128, momentum=momentum)
9
10 self.fc5 = nn.Linear(midfeat, 768, bias=False)
11 self.bn5 = nn.BatchNorm1d(768, momentum=momentum)
12
13 self.gru = nn.GRU(input_size=768, hidden_size=256, num_layers=2,
14 bias=True, batch_first=True, dropout=0., bidirectional=True)
15
16 self.fc = nn.Linear(512, classes_num, bias=True)
17
18 self.init_weight()Source: pytorch/models.py
设计意图：
- 卷积层负责局部声学模式（谐波、起音瞬态），GRU 负责时序上下文（延音、踏板造成的长时能量包络），这是转写任务的标准 CRNN 范式。
- 权重初始化被显式定制：
init_layer使用 Xavier 均匀初始化并置零偏置，init_gru对 GRU 的输入-隐藏权重使用均匀分布、隐藏-隐藏权重对其中 1/3 使用正交初始化——正交初始化有助于 RNN 训练稳定性：
python
1def init_gru(rnn):
2 """Initialize a GRU layer. """
3
4 def _concat_init(tensor, init_funcs):
5 (length, fan_out) = tensor.shape
6 fan_in = length // len(init_funcs)
7
8 for (i, init_func) in enumerate(init_funcs):
9 init_func(tensor[i * fan_in : (i + 1) * fan_in, :])
10
11 def _inner_uniform(tensor):
12 fan_in = nn.init._calculate_correct_fan(tensor, 'fan_in')
13 nn.init.uniform_(tensor, -math.sqrt(3 / fan_in), math.sqrt(3 / fan_in))
14
15 for i in range(rnn.num_layers):
16 _concat_init(
17 getattr(rnn, 'weight_ih_l{}'.format(i)),
18 [_inner_uniform, _inner_uniform, _inner_uniform]
19 )
20 torch.nn.init.constant_(getattr(rnn, 'bias_ih_l{}'.format(i)), 0)
21
22 _concat_init(
23 getattr(rnn, 'weight_hh_l{}'.format(i)),
24 [_inner_uniform, _inner_uniform, nn.init.orthogonal_]
25 )
26 torch.nn.init.constant_(getattr(rnn, 'bias_hh_l{}'.format(i)), 0)Source: pytorch/models.py
classes_num 由顶层模型（如 Regress_onset_offset_frame_velocity_CRNN）决定：音符模型输出 onset/offset/frame/velocity 等多组概率图，踏板模型输出踏板相关目标，最终在推理端经阈值化与回归时间映射为 MIDI 事件。
端到端工作流（Core Flow）
runme.sh 是理解整个系统如何运转的最佳入口，它按顺序执行：下载预训练模型并推理 → 配置数据集 → 打包 hdf5 → 训练音符模型 → 训练踏板模型 → 合并权重 → 评测。关键命令如下（摘自 runme.sh 实际内容）：
bash
# --- 1. Train note transcription system ---
python3 pytorch/main.py train --workspace=$WORKSPACE --model_type='Regress_onset_offset_frame_velocity_CRNN' --loss_type='regress_onset_offset_frame_velocity_bce' --augmentation='none' --max_note_shift=0 --batch_size=12 --learning_rate=5e-4 --reduce_iteration=10000 --resume_iteration=0 --early_stop=300000 --cudaSource: runme.sh
bash
# --- 2. Train pedal transcription system ---
python3 pytorch/main.py train --workspace=$WORKSPACE --model_type='Regress_pedal_CRNN' --loss_type='regress_pedal_bce' --augmentation='none' --max_note_shift=0 --batch_size=12 --learning_rate=5e-4 --reduce_iteration=10000 --resume_iteration=0 --early_stop=300000 --cudaSource: runme.sh
bash
1# --- 3. Combine the note and pedal models ---
2NOTE_CHECKPOINT_PATH="Regress_onset_offset_frame_velocity_CRNN_onset_F1=0.9677.pth"
3PEDAL_CHECKPOINT_PATH="Regress_pedal_CRNN_onset_F1=0.9186.pth"
4NOTE_PEDAL_CHECKPOINT_PATH="CRNN_note_F1=0.9677_pedal_F1=0.9186.pth"
5python3 pytorch/combine_note_and_pedal_models.py --note_checkpoint_path=$NOTE_CHECKPOINT_PATH --pedal_checkpoint_path=$PEDAL_CHECKPOINT_PATH --output_checkpoint_path=$NOTE_PEDAL_CHECKPOINT_PATHSource: runme.sh
评测分两步进行：先在测试集上推理概率，再离线计算指标：
bash
1# Inference probability for evaluation
2python3 pytorch/calculate_score_for_paper.py infer_prob --workspace=$WORKSPACE --model_type='Note_pedal' --checkpoint_path=$NOTE_PEDAL_CHECKPOINT_PATH --augmentation='none' --dataset='maestro' --split='test' --cuda
3
4# Calculate metrics
5python3 pytorch/calculate_score_for_paper.py calculate_metrics --workspace=$WORKSPACE --model_type='Note_pedal' --augmentation='aug' --dataset='maestro' --split='test'
6python3 pytorch/calculate_score_for_paper.py calculate_metrics --workspace=$WORKSPACE --model_type='Note_pedal' --augmentation='aug' --dataset='maps' --split='test'Source: runme.sh
本地快速推理（下载官方 Zenodo checkpoint 后）：
bash
1CHECKPOINT_PATH="CRNN_note_F1=0.9677_pedal_F1=0.9186.pth"
2wget -O $CHECKPOINT_PATH "https://zenodo.org/record/4034264/files/CRNN_note_F1%3D0.9677_pedal_F1%3D0.9186.pth?download=1"
3MODEL_TYPE="Note_pedal"
4python3 pytorch/inference.py --model_type=$MODEL_TYPE --checkpoint_path=$CHECKPOINT_PATH --audio_path='resources/cut_liszt.mp3' --cudaSource: runme.sh
整个"训练 → 合并 → 推理"的时间线交互如下：
Loading diagram...
训练数据与基线结果
训练使用 MAESTRO V2.0.0 数据集（约 200 小时名家钢琴演奏，音频与 MIDI 标签对齐精度约 3 ms）。README 给出的统计如下：
| Split | Performances | Duration (hours) | Size (GB) | Notes (millions) |
|---|---|---|---|---|
| Train | 967 | 161.3 | 97.7 | 5.73 |
| Validation | 137 | 19.4 | 11.8 | 0.64 |
| Test | 178 | 20.5 | 12.4 | 0.76 |
| Total | 1282 | 201.2 | 121.8 | 7.13 |
Source: README.md
在单张 Tesla-V100-PCIE-32GB 上训练 300k 迭代（约一周），README 中记录的训练日志（Namespace 一行）展示了 main.py 的完整超参数集合：
text
Namespace(augmentation='none', batch_size=12, cuda=True, early_stop=300000, filename='main', learning_rate=0.0005, loss_type='regress_onset_offset_frame_velocity_bce', max_note_shift=0, mini_data=False, mode='train', model_type='Regress_onset_offset_frame_velocity_CRNN', reduce_iteration=10000, resume_iteration=0, workspace='.../workspaces/piano_transcription')Source: README.md
训练收敛后的验证/测试统计（第 300000 次迭代）显示 frame_ap ≈ 0.93，回归目标（onset/offset/velocity MAE）均收敛到小数值：
text
1Iteration: 300000
2 Train statistics: {'frame_ap': 0.9439, 'reg_onset_mae': 0.091, 'reg_offset_mae': 0.127, 'velocity_mae': 0.0241}
3 Validation statistics: {'frame_ap': 0.9245, 'reg_onset_mae': 0.0985, 'reg_offset_mae': 0.1327, 'velocity_mae': 0.0265}
4 Test statistics: {'frame_ap': 0.9285, 'reg_onset_mae': 0.097, 'reg_offset_mae': 0.1353, 'velocity_mae': 0.027}Source: README.md
checkpoint 与统计量按 workspace/statistics/{filename}/{model_type}/loss_type=.../augmentation=.../batch_size=.../ 的目录层级落盘（见 README 训练日志中的路径样例），便于超参数组合之间的横向比较。
使用示例
基础用法：pip 包推理（推荐）
最简单的转写方式是安装官方推理包 piano_transcription_inference（无需本仓库训练代码）：
python
1from piano_transcription_inference import PianoTranscription, sample_rate, load_audio
2
3# Load audio
4(audio, _) = load_audio('resources/cut_liszt.mp3', sr=sample_rate, mono=True)
5
6# Transcriptor
7transcriptor = PianoTranscription(device='cuda') # 'cuda' | 'cpu'
8
9# Transcribe and write out to MIDI file
10transcribed_dict = transcriptor.transcribe(audio, 'cut_liszt.mid')Source: README.md
高级用法：Replicate/Cog 在线推理接口
仓库根目录的 predict.py 是面向 Replicate 平台的 Cog 预测器：setup() 中用本地 model.pth 构建 PianoTranscription，predict() 中加载任意上传音频、转写为 MIDI，并用 synthviz 生成可视化视频返回：
python
1class Predictor(cog.Predictor):
2 transcriptor: PianoTranscription
3
4 def setup(self):
5 self.transcriptor = PianoTranscription(
6 device="cuda", checkpoint_path="./model.pth"
7 )
8
9 @cog.input("audio_input", type=Path, help="Input audio file")
10 def predict(self, audio_input):
11 midi_intermediate_filename = "transcription.mid"
12 video_filename = os.path.join(Path.cwd(), "output.mp4")
13 audio, _ = librosa.core.load(str(audio_input), sr=sample_rate)
14 # Transcribe audio
15 self.transcriptor.transcribe(audio, midi_intermediate_filename)
16
17 # 'Visualization' output option
18 create_video(
19 input_midi=midi_intermediate_filename, video_filename=video_filename
20 )
21 print(
22 f"Created video of size {os.path.getsize(video_filename)} bytes at path {video_filename}"
23 )
24 # Return path to video
25 return Path(video_filename)Source: predict.py
这一封装展示了系统的最小可用形态：加载音频（重采样到 sample_rate）→ PianoTranscription.transcribe(audio, midi_path) → 得到 MIDI（及可视化）。
配置选项（main.py 训练命令行参数）
从 runme.sh 与 README 训练日志可确认的 pytorch/main.py 训练参数：
| 参数 | 类型 | 默认/示例值 | 说明 |
|---|---|---|---|
mode | string | train | 运行模式（train） |
workspace | string | .../workspaces/piano_transcription | 工作区根目录，存放 hdf5/checkpoints/statistics |
model_type | string | Regress_onset_offset_frame_velocity_CRNN / Regress_pedal_CRNN / Note_pedal | 模型类型：音符系统 / 踏板系统 / 合并后的推理模型 |
loss_type | string | regress_onset_offset_frame_velocity_bce / regress_pedal_bce | 损失函数类型，与 model_type 配套 |
augmentation | string | none | 数据增强开关 |
max_note_shift | int | 0 | 训练时最大音高偏移（数据增强） |
batch_size | int | 12 | 批大小；约需 29 GB 显存，OOM 时需调小 |
learning_rate | float | 5e-4 | 学习率 |
reduce_iteration | int | 10000 | 学习率衰减间隔 |
resume_iteration | int | 0 | 从第 N 次迭代恢复训练 |
early_stop | int | 300000 | 总迭代数上限（提前停止） |
cuda | flag | true | 是否使用 GPU |
filename | string | main | 实验名，用于统计目录命名 |
mini_data | bool | False | 是否使用小数据子集（调试用） |
推理相关参数（pytorch/inference.py）：model_type、checkpoint_path、audio_path、cuda；评测参数（calculate_score_for_paper.py）：mode（infer_prob / calculate_metrics）、workspace、model_type、checkpoint_path、augmentation、dataset（maestro / maps）、split、cuda。
专业说明（性能、边界与运维）
- 显存与批量：README 明确指出 batch size 12 需要约 29 GB GPU 显存，训练使用单张 Tesla-V100-PCIE-32GB；FAQ 建议"遇到 GPU 内存不足时降低 batch size"，也可以考虑多卡训练。
- 训练时长：300k 迭代约一周（单卡 V100），README 日志中第 300000 迭代一轮 train + validate 用时约 8953 s + 94 s。
- 数据准备前置条件：训练前必须先下载 MAESTRO V2.0.0 并按
runme.sh打包为 hdf5；数据集目录结构见 README（按年份 2004–2018 分目录 +maestro-v2.0.0.csv/json）。 - 版本兼容性：基于 Python 3.7 + PyTorch 1.4.0 开发，官方声明"其他版本应可用但未完整测试"，升级依赖时需注意。
- 推理路径分离：本仓库面向研究与从头训练；生产化推理走
piano_transcription_inference包或 Replicate/Cog（predict.py）。 - 许可证与引用：Apache 2.0；引用论文 [1] Kong et al., "High-resolution Piano Transcription with Pedals by Regressing Onsets and Offsets Times", arXiv:2010.01815。
- 下游应用：基于本系统构建了 GiantMIDI-Piano 大规模古典钢琴 MIDI 数据集（README "Applications"）。
相关链接（Related Links）
- 论文：High-resolution Piano Transcription with Pedals by Regressing Onsets and Offsets Times (arXiv:2010.01815)
- 推理包：piano_transcription_inference
- 下游数据集：GiantMIDI-Piano
- 在线 Demo：Replicate - bytedance/piano-transcription
- 关键源码入口：runme.sh、pytorch/main.py、pytorch/models.py、pytorch/inference.py、pytorch/combine_note_and_pedal_models.py、predict.py
Sources
(3 files)(root)
pytorch