fairseq S2ST 模型基准测试框架指南运行时、内存与 FLOPS 测量实战【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm导读本文讲解 unilm 仓库中 kosmos-2/fairseq 子项目提供的 Speech-to-Speech TranslationS2ST模型基准测试框架benchmarking 目录。该框架用于在统一条件下量化各类 S2ST 模型级联模型与端到端模型的推理运行时runtime、最大内存占用max memory与浮点运算量FLOPS并支持任意 fairseq 模型扩展。读完本文你将掌握该框架的三项核心指标定义、完整 CLI 调用方式、npy/raw 两类数据集格式以及其底层Processing、SpeechGeneration、S2UT、Cascaded2StageS2ST、Cascaded3StageS2ST五类基准测试实现的工作原理。框架目标与设计原则S2ST 任务通常存在两条技术路线端到端模型如直接语音到离散单元翻译 S2UT与级联模型如 ST TTS 两阶段、ASR MT TTS 三阶段。两类模型在结构上差异巨大若直接比较整体延迟会混入数据中间转换、IO 等与模型本身无关的开销导致不公平。该框架对此的处理方式是对于级联模型只统计各阶段模型推理本身的指标忽略所有中间数据传递与 IO 处理耗时。在代码层面这一原则体现在各级联类的compute_metrics方法中它先把上游输出转换为下游输入如generate_s2t_outputs、generate_tts_inputs再分别对各阶段单独测指标最后按规则聚合详见下文级联模型的指标聚合。同时框架明确在CPU上执行所有基准测试理由是CPU 是生产环境中常见的部署形态且缺少对 GPU 友好的成熟基准测试库支持。这一点也体现在默认配置cpu: True以及 get_metrics.py 中config[general]对命令行参数的覆盖上。为了保证结果可复现入口脚本 get_metrics.py 与 core.py 都会执行torch.manual_seed(1)与torch.set_deterministic(True)data_utils.py 同样对random、np.random设置种子。三项核心指标定义框架围绕三类指标设计每个指标都有一套独立的测量方法指标定义单位测量手段Runtime对数据集中单个样本执行一次模型推理的平均耗时秒secPython 标准库timeitMax memory对数据集中全部样本执行推理期间单样本最大内存占用的平均值MiBmemory_profiler的memory_usageFLOPS对数据集中单个样本执行推理所需的平均浮点运算次数次flopsPAPI 库的PAPI_DP_OPS计数器三者均在 core.py 的gather_all_metrics中按顺序调用并一次返回def gather_all_metrics(self, dataset, repeat): run_time self.benchmark_run_time(dataset, repeat) max_memory self.max_memory(dataset, repeat) flops self.count_flops(dataset, repeat) return run_time, max_memory, flopsRuntime基于 timeit 的平均单样本推理耗时在 core.py 中单样本耗时通过timeit.Timer(lambda: self.forward(sample))构造计时器timer.timeit(repeat)执行多次后取平均def benchmark_run_time_single_sample(self, sample, repeat): timer timeit.Timer(lambda: self.forward(sample)) time_elapsed timer.timeit(repeat) return time_elapsed / repeat外层benchmark_run_time会遍历整个数据集将每个样本的平均耗时累加后除以样本数得到数据集层面的平均推理时间每处理 100 个样本会打印一次进度日志。Max memory基于 memory_profiler 的峰值内存core.py 使用memory_usage((self.forward, (sample,), {}))捕获执行forward(sample)期间的内存曲线并取最大值total_memory max(memory_usage((self.forward, (sample,), {})))最后除以repeat * len(dataset)得到平均最大内存。对级联模型取各阶段最大内存的最大值作为整体峰值因为阶段间通常串行执行峰值应取 max 而非求和。FLOPS基于 PAPI 的浮点运算计数core.py 借助pypapi的硬件计数器统计浮点双精度运算次数high.start_counters([events.PAPI_DP_OPS]) for i, sample in enumerate(dataset): for _r in range(repeat): self.forward(sample) flops high.stop_counters() flops round(flops[0] / (repeat * len(dataset)))注释中明确说明PAPI 计数器只在模型运行于 CPU 时有效这再次呼应了框架全部在 CPU 上测的前提约束。使用前需要安装pypapi依赖并确保系统支持对应硬件计数器。CLI 命令与入口脚本原文档给出的核心命令如下注意get_metrics.py后需要传入空字符串占位参数再以--config指定 YAML 配置文件CUBLAS_WORKSPACE_CONFIG:4096:8 python examples/speech_to_speech/benchmarking/get_metrics.py --config $configCUBLAS_WORKSPACE_CONFIG:4096:8是 cuBLAS 的工作区配置用于保证 GPU 上矩阵运算结果的确定性与代码中torch.set_deterministic(True)相互配合。入口脚本 get_metrics.py 的make_parser定义了完整的命令行参数参数默认值可选值说明--config无任意 YAML 路径基准测试配置文件内容会合并进命令行参数--model-typeS2US2S、TTS、S2UT、MT、S2T、2StageS2ST、3StageS2ST选择要测试的模型类型对应 core.py 中不同的实现类--dataset-path无文件路径测试数据集路径npy 或 raw--dataset-typenpynpy、raw数据集文件类型--read-using-sfFalse布尔raw 数据集是否用 soundfile 读取--dataset-sizeNone整数参与测试的样本数量上限便于快速小规模验证--dump-speech-waveforms-dirNone目录路径若指定则将模型输出的语音波形 dump 到该目录--dump-waveform-file-prefix字符串导出的 wav 文件名前缀--feat-dim80整数输入特征维度--target-sr16000整数dump 波形时的目标采样率此外解析器还挂载了 fairseq 的 generation 参数options.add_generation_args与交互式生成参数options.get_interactive_generation_parser用于透传 beam size、max_len_a/b等推理配置。配置文件的分阶段加载机制cli_main中的核心逻辑是将 YAML 配置按general/stage1/stage2/stage3分片合并进命令行参数with open(args.config, r) as f: config yaml.load(f, Loaderyaml.FullLoader) dict_args vars(args) dict_args.update(config[general]) # general 段覆盖全局默认参数 i 1 stage_args [] while i 3: var fstage{i} tmp_args copy.deepcopy(dict_args) if var in config: tmp_args.update(config[var]) # 每阶段配置覆盖对应模型参数 stage_args.append(Namespace(**tmp_args)) i 1 else: break随后根据model_type将stage_args分发给不同的类if args.model_type in (S2S, TTS): model SpeechGeneration(stage_args[0]) elif args.model_type S2UT: model S2UT(stage_args[0], stage_args[1] if len(stage_args) 1 else None) elif args.model_type in (MT, S2T): model Processing(stage_args[0]) elif args.model_type 2StageS2ST: model Cascaded2StageS2ST(stage_args[0], stage_args[1]) elif args.model_type 3StageS2ST: model Cascaded3StageS2ST(stage_args[0], stage_args[2], stage_args[1])命名约定见make_parser的 docstrings2x_args用于语音输入模型ST、ASRx2s_args用于语音输出模型TTSmt_args用于翻译模型MT、T2U直接式 S2ST 模型则用x2s_args提供模型信息。基准测试的主流程为按dataset_type加载数据集model.warm_up(sampledataset[0], repeat2)用首个样本预热模型两次core.pymodel.gather_all_metrics(dataset, repeat1)一次采集三类指标并打印run_time / memory / flops若指定了--dump-speech-waveforms-dir则通过dump_final_speech_output将每个样本的模型输出写为{prefix}_{i}_pred.wavcore.py便于人工核对模型是否真的生成了有效波形。数据集格式npy 与 raw原文档说明框架支持两类数据集并在 data_utils.py 中给出了对应加载实现。npy 数据集npy 数据集是一个以.npy文件保存的样本列表每个样本是一个字典结构如下sample: { id: xx, net_input: { src_tokens: torch.tensor([]), src_lengths: torch.tensor([]) } }其中src_tokens为输入 token对语音输入通常是音频特征张量src_lengths为对应长度。加载函数为load_dataset_npy使用np.load(file_name, allow_pickleTrue)读取并可通过dataset_size截取前 N 个样本。该格式可直接由 data_utils.py 的cli_main生成给定--data、--subset、--task先构建 task 并加载子集然后分别生成三类子集并保存random_dataset{token}w_ids.npy随机抽取 500 条样本short_dataset{token}w_ids.npy按src_lengths升序排序后取前 500 条get_short_data_subsetlong_dataset{token}w_ids.npy按src_lengths降序排序后取前 500 条get_long_data_subset。同时还支持--ref-dataset按 id 过滤、--dataset-save-token给文件名加标记。生成后会打印数据集的max_len / min_len / avg_len统计信息get_dataset_stats。raw 数据集raw 数据集是一个类似 wav2vec2 输入 tsv 的音频路径清单文件首行是音频根目录其后每行第一列以\t分隔为相对路径。load_dataset_raw_to_waveforms会将其解析为完整路径并默认按 16 kHz 采样率读取波形若read_using_soundfileTrue用sf.read读取双声道取均值.mean(-1)否则走 fairseq 的get_waveform路径把波形转为{id, net_input: {src_tokens, src_lengths}}结构。此外data_utils.py还提供了generate_random_dataset/generate_random_data_sample默认特征维度 D80可快速构造随机测试集用于框架自检。五类基准测试实现从源码理解测量原理core.py 以BenchmarkingBase(nn.Module)为基类定义了warm_up、benchmark_run_time、count_flops、max_memory、gather_all_metrics、dump_final_speech_output等通用方法再派生出五类具体实现。ProcessingASR / MT / ST 模型Processing复刻了fairseq_cli/generate.py的推理流程core.py通过tasks.setup_task构建任务、checkpoint_utils.load_model_ensemble加载模型暂不支持多模型 ensemble、task.build_generator构建生成器并处理 tokenizer / BPE /post_process。encode_source把文本编码为net_inputforward调用task.inference_step得到假设decode_target把 token 解码回文本。它对应--model-type中的S2T与MT。SpeechGenerationTTS 与直接式 S2STSpeechGeneration复刻了examples/text_to_speech/generate_waveform.py的流程core.py面向以语音生成为目标的模型TTS、直接式 S2ST 等。注意setUp中当args.task speech_to_speech时会把normalize_waveform置为False。processTextInput将文本编码为带prev_output_tokens的net_inputforward直接调用generator.generate。对应--model-type的S2S与TTS。GenerateWaveformFromCode 与 HubertUnitExtractor辅助模块GenerateWaveformFromCodecore.py使用CodeHiFiGANVocoder从离散单元code生成波形支持dur_prediction时长预测开关对应--target-is-code场景。HubertUnitExtractorcore.py基于HubertFeatureReader与ApplyKmeans提取 HuBERT 特征并量化为 k-means 单元用于构造单元类模型的输入。S2UT端到端语音到单元翻译S2UTcore.py内部组合Processings2u 阶段与可选的GenerateWaveformFromCodevocoder 阶段。其forward先做 s2u 推理并解码为单元串再送入 vocoder 生成波形。指标计算通过compute_metrics完成vocoder 的输入vocoder_input由generate_s2u_outputs预生成一次随后两阶段分别测同一种指标内存取两阶段最大值运行时与 FLOPS 取两阶段之和。Cascaded2StageS2STST TTSCascaded2StageS2STcore.py组合Processings2t与SpeechGenerationtts。forward依次执行语音识别/翻译得到文本再processTextInput送入 TTS。指标聚合方式与 S2UT 一致先预生成tts_inputs即把数据集整体跑一遍 s2t 并转成 tts 输入再分别测 s2t 与 tts内存取 max、其余取和从而完全剥离中间文本传递与 IO 的开销——这正是原文档宣称级联模型只统计模型推理指标的实现落点。Cascaded3StageS2STASR MT TTSCascaded3StageS2STcore.py继承两阶段类并增加Processingmt作为中间翻译阶段。compute_metrics中先预生成mt_inputs与tts_inputs再分别对 s2t、mt、tts 三段测指标打印三者明细内存取三者最大值运行时与 FLOPS 三者求和。配置文件详解仓库内置四份示例配置统一采用generalstageN结构S2TS2T.yamlgeneral: dataset_path: $npy_dataset cpu: True model_type: S2T dataset_size: 1 stage1: data: $data_bin task: speech_to_text path: $checkpoint config_yaml: config.yaml max_len_a: 2 max_len_b: 5002StageS2ST2StageS2ST.yamlgeneral: dataset_path: $npy_dataset cpu: True model_type: 2StageS2ST dataset_size: 1 stage1: data: $data_bin_stage1 task: speech_to_text path: $checkpoint_stage1 config_yaml: config.yaml max_len_a: 2 max_len_b: 500 stage2: data: $data_bin_stage2 task: text_to_speech path: $checkpoint_stage2 config_yaml: config.yaml3StageS2ST3StageS2ST.yamlgeneral: dataset_path: $npy_dataset cpu: True model_type: 3StageS2ST max_len_a: 2 max_len_b: 500 dataset_size: 1 stage1: data: $data_bin_stage1 task: speech_to_text path: $checkpoint_stage1 config_yaml: config.yaml max_len_a: 2 max_len_b: 500 stage2: data: $data_bin_stage2 task: translation path: $checkpoint_stage2 config_yaml: config.yaml stage3: data: $data_bin_stage3 task: text_to_speech path: $checkpoint_stage3 config_yaml: config.yaml需要提醒仓库中的 3StageS2ST.yaml 存在两个重复的stage2键YAML 规范下后一个会覆盖前一个。由于get_metrics.py对三阶段模型按stage1 → stage2(translation) → stage3(text_to_speech)读取Cascaded3StageS2ST(stage_args[0], stage_args[2], stage_args[1])即 s2t、mt、tts 分别对应 stage1、stage2、stage3实际使用时请按上面的结构确保stage3为text_to_speech阶段避免配置被覆盖导致阶段缺失。DirectS2UDirectS2U.yamlgeneral: dataset_path: $npy_dataset_path cpu: True model_type: S2UT dataset_size: 5 dump_speech_waveforms_dir: $dump_waveforms_dir_path stage1: data: $data_bin task: speech_to_speech path: $checkpoint config_yaml: config.yaml max_len_b: 100000 beam: 10 target_is_code: True max_target_positions: 3000 target_code_size: 100 stage2: vocoder: $vocoder_path vocoder_cfg: $vocoder_cfg_json dur_prediction: True该配置演示了 S2UT 的典型设置target_is_code: True表示目标为离散单元对应--target-is-code参数target_code_size: 100与 S2UT 常用的 100 单元 HuBERT k-means 字典一致相关背景见 speech_to_speech README 中关于 HuBERT Base layer 6 与 100 units 的说明stage2给出 vocoder 权重、配置 JSON 与dur_prediction开关。使用步骤与注意事项准备数据集使用 data_utils.py 生成 npy 数据集随机/短/长三份或准备 wav2vec2 风格的 tsv 路径清单作为 raw 数据集。编写配置文件按模型类型从configs/下四份示例中选取并复制一份替换$开头的占位变量dataset_path、data、path、vocoder、vocoder_cfg等并确认task、model_type、dataset_size正确。运行基准测试CUBLAS_WORKSPACE_CONFIG:4096:8 python examples/speech_to_speech/benchmarking/get_metrics.py --config $config脚本会先打印数据集路径随后输出类似run_time Xsec memory YMiB flops Z的结果。验证输出可选配置dump_speech_waveforms_dir与--dump-waveform-file-prefix后模型生成的波形会以 wav 形式落盘便于检查推理产物是否有效。注意事项框架全部在 CPU 上测量FLOPS指标依赖 PAPI 硬件计数器需安装pypapi且平台支持PAPI_DP_OPS内存与 FLOPS 测量会遍历整个数据集dataset_size可控制样本量以缩短测试时间多模型 ensemble 目前不受支持len(models) 1会直接抛异常需要安装memory_profiler、pypapi、soundfile等第三方依赖具体依赖可参考 fairseq 安装说明 与项目 requirements。小结该基准测试框架以公平、可复现、易扩展为设计目标通过剥离级联模型的中间开销保证端到端与级联模型可比通过固定随机种子与确定性设置保证结果可复现通过BenchmarkingBase基类 分阶段 YAML 配置保证新模型接入成本极低。核心实现集中在 core.py五类模型实现与三项指标测量、get_metrics.pyCLI 与配置合并、data_utils.py数据加载与生成三个文件中配合configs/下四份示例配置即可快速开展 S2ST 模型的运行时、内存与 FLOPS 基准测试。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考