8.2.5. Example: Speech Recognition using Whisper
Whisper の base.en モデルを使用し、LibriSpeech ASR corpus の test-clean データセットで音声認識の推論を行うサンプルプログラムです。この例では、認識したテキストの単語誤り率（WER）を出力します。
注釈
このサンプルの Whisper 由来のソースコード、Whisper の実装、および学習済みの重みは、MIT License の下で提供されています。
LibriSpeech ASR corpus は CC BY 4.0 の下で提供されています。
ライセンスとソースコードの由来
配布するサンプルには、openai/whisper のコードをコピー・改変した部分と、独自に再実装した部分が含まれます。データセットのラッパー、推論の流れ、出力の後処理、および static_whisper.py のモデルの forward と attention の処理が対象です。コードの由来の詳細は、サンプルと同じディレクトリの license_note.md を参照してください。
OpenAI の著作権表示と MIT ライセンス全文は、サンプルと同じディレクトリの LICENSE.whisper に収録しています。
以下の上流ソースの参照先はコミット c0d2f624c09dc18e709e37c2ad90c039a4eb72a2 に固定しています。モデル部分には、独自に再実装した処理と、上流の forward や attention のコードをコピー・改変した処理が含まれます。MLSDK 向けにモデルの分割、KV キャッシュ、および forward の差し替えを変更しています。
notebooks/LibriSpeech.ipynb: データセットのラッパー、推論の流れ、および Word Error Rate の算出処理。
whisper/model.py:
Whisper、AudioEncoder、TextDecoder、ResidualAttentionBlock、MultiHeadAttention。AudioEncoder.forward、TextDecoder.forward、ResidualAttentionBlock.forward、MultiHeadAttention.forward／qkv_attentionの処理を含みます。whisper/decoding.py:
DecodingTask、PyTorchInference、decode()。デコーディングループ、候補の選択、およびテキストの後処理で使用・参照しています。
このサンプルは Whisper パッケージと学習済みの重みを実行時に取得します。モデル内部の forward を差し替えるため、openai-whisper は参照元ソースに記載されたバージョン 20250625 に固定しています。LibriSpeech も実行時に取得します。SDK が配布するのは改変したサンプルのソースであり、データセットには CC BY 4.0 が適用されます。
Preparation
この例では LibriSpeech データセットを使用します。初回実行時に、torchaudio は選択したデータセット片を ./librispeech_dataset に自動でダウンロードして展開します。デフォルトの分割は test-clean です。
たとえばパーミッションエラーにより自動ダウンロードに失敗した場合は、データセットを手動でダウンロードして展開してください。データセットのダウンロードリンクは LibriSpeech ASR corpus のページを参照してください。
$ mkdir librispeech_dataset
$ curl -L https://us.openslr.org/resources/12/test-clean.tar.gz -o librispeech_dataset/test-clean.tar.gz
$ tar zxvf librispeech_dataset/test-clean.tar.gz --no-same-owner -C librispeech_dataset/
Execution Method
このサンプルのディレクトリから次のコマンドを実行します。
$ cd /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference
$ ./run_whisper_inference.sh [OPTION]...
たとえば、次のコマンドは MN-Core 2 で音声認識の推論を実行します。
$ ./run_whisper_inference.sh -b mncore2:auto
最初の 10 イテレーションだけを処理する簡易動作確認を行うには、次のように実行します。
$ ./run_whisper_inference.sh -b mncore2:auto --max_iterations 10
Command-line Options
以下のオプションを使用できます。configs.toml で指定されるオプションもコマンドラインで渡すことができます。
-h、--help: 使用可能なコマンドラインオプションを表示します。--dataset_dir DATASET_DIR: LibriSpeech データセットをダウンロードして保存するディレクトリを指定します。デフォルトは、サンプルディレクトリ内の./librispeech_datasetです。--data_name DATA_NAME: 使用する LibriSpeech データセットの分割を指定します。デフォルトはtest-cleanです。--chunk_length CHUNK_LENGTH: 各入力音声チャンクの長さを秒単位で指定します。デフォルトは30です。--model MODEL: 読み込む Whisper モデルを指定します。デフォルトはbase.enです。--model_dir MODEL_DIR: Whisper モデルファイルをダウンロードして読み込むディレクトリを指定します。デフォルトは/tmp/whisper_modelsです。--seed SEED: 実行を再現可能にするための乱数シードを指定します。デフォルトは0です。--batch_size BATCH_SIZE: 各バッチで処理する音声サンプル数を指定します。デフォルトは1です。--max_iterations MAX_ITERATIONS: DataLoader の最大イテレーション数を指定します。デフォルト値の-1はデータセット全体を処理します。--wer_criteria WER_CRITERIA: 許容する Word Error Rate (WER) の最大値を指定します。デフォルトは0.05(5%) です。測定された WER がこの値を超えると、実行は失敗します。-b BACKEND、--backend BACKEND: 推論に使用するバックエンドを指定します。サポートされる値はcpu、pfvm:cpu、mncore2:auto、mncore2:[0-7]、emu2です。デフォルトはcpuです。--out_basedir OUT_BASEDIR: 生成される出力成果物のベースディレクトリを指定します。デフォルトは/tmpです。--optimize_option {debug,O0,O1,O2,O3,O4}:Context.compile()に渡す最適化オプションを指定します。デフォルトはO1です。--preset_options_dir PRESET_OPTIONS_DIR: プリセット最適化オプションの JSON ファイルを含むディレクトリを指定します。デフォルトは/opt/pfn/pfcomp/codegen/preset_optionsです。ラッパースクリプトはこれを Codegen ソースツリーのpreset_optionsディレクトリで上書きします。--disable_compile_cache: コンパイル済み成果物の再利用を無効にします。デフォルトでは、バックエンドが MLSDK を使用する場合にコンパイル済み成果物が再利用されます。--trace_filename TRACE_FILENAME: トレースを有効にする出力ファイル名を指定します。デフォルトではトレースは無効です。
Expected Output
プログラムはいくつかの認識結果、推論の性能情報と、認識テキストおよび LibriSpeech の正解文字起こしから計算した WER を出力します。例:
############ Outputs #############
Normalized references and hypotheses (first 3 samples):
[0]
Reference : he hoped there would be stew for dinner turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick peppered flour fattened sauce
Hypothesis: he hoped there would be stew for dinner turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick peppered flower fat and sauce
[1]
Reference : stuff it into you his belly counseled him
Hypothesis: stuffered into you his belly counseled him
[2]
Reference : after early nightfall the yellow lamps would light up here and there the squalid quarter of the brothels
Hypothesis: after early nightfall the yellow lamps would light up here and there the squalid quarter of the brothels
##################################
########## Performances ##########
---------- configs ----------
sample name: whisper_inference
backend device: mncore2:auto
onnx exporter: fx2onnx
optimizer: None
optimize_option: O1
------------------------------
---------- performance of evaluation/inference part ----------
eval epochs: 1
batch size: 1
iterations per epoch: 100
total time [s]: 14.785252013942227
average per epoch [s]: 14.785252013942227
averaged [iter/s]: 6.763496483232196
--------------------------------------------------------------
##################################
Result: WordErrorRate(WER) = 2.55 %
デフォルトでは、WER が 5% を超えると実行は失敗します。この閾値は --wer_criteria で変更できます。
Scripts
1#! /bin/bash
2
3set -eux -o pipefail
4
5CURRENT_DIR=$(realpath $(dirname $0))
6
7EXAMPLE_NAME="whisper_inference"
8VENV_DIR=/tmp/${EXAMPLE_NAME}_venv
9
10# Whisper requires ffmpeg for audio processing.
11if ! command -v ffmpeg > /dev/null 2>&1; then
12 apt update && apt install -y ffmpeg
13fi
14
15if [[ ! -d ${VENV_DIR} ]]; then
16 python3 -m venv --system-site-packages ${VENV_DIR}
17 source ${VENV_DIR}/bin/activate
18 pip install -r ${CURRENT_DIR}/requirements.txt
19 pip install -r ${CURRENT_DIR}/requirements_cpu.txt --index-url https://download.pytorch.org/whl/cpu
20else
21 source ${VENV_DIR}/bin/activate
22fi
23
24CODEGEN_DIR=${CODEGEN_DIR:-${CURRENT_DIR}/../../../}
25BUILD_DIR=${BUILD_DIR:-${CODEGEN_DIR}/build}
26source "${BUILD_DIR}/codegen_pythonpath.sh"
27
28# To enable deterministic behavior, set env. var. as follows
29export CUBLAS_WORKSPACE_CONFIG=:16:8
30
31# MLSDK config
32export CODEGEN_GEMM_FORCE_WEIGHT_ON_DRAM=1
33
34exec python3 ${CURRENT_DIR}/${EXAMPLE_NAME}.py "$@" \
35 --preset_options_dir ${CODEGEN_DIR}/preset_options
1# Portions adapted from openai/whisper. Copyright (c) 2022 OpenAI.
2# See LICENSE.whisper for the MIT License and license_note.md for source references.
3
4import argparse
5import os
6from collections.abc import Callable
7from itertools import islice
8
9import numpy as np
10import torch
11import whisper
12from mlsdk import (
13 CompiledFunction,
14 Context,
15 MNDevice,
16 TensorLike,
17 TensorProxy,
18 trace_scope,
19)
20from static_whisper import (
21 AudioEncoderHead,
22 AudioEncoderTransformer,
23 CrossKVPrecompute,
24 TextDecoderEmbedding,
25 TextDecoderTransformer,
26)
27from tqdm import tqdm
28from utility import (
29 DeviceSet,
30 Timer,
31 apply_toml_defaults,
32 compile_fn,
33 decide_outdir,
34 output_result_times,
35 register_model,
36 set_deterministic_mode,
37)
38from whisper.decoding import DecodingOptions, DecodingTask
39from whisper_utils import (
40 LibriSpeech,
41 calc_wer,
42 kv_cache_from_dict,
43 prepare_cache_sample,
44)
45
46DEVICES = DeviceSet(["pfvm:cpu", "mncore2:auto", "mncore2:[0-7]", "emu2", "cpu"])
47MNCORE_DEVICES = DeviceSet(["mncore2:auto", "mncore2:[0-7]"])
48PFVM_DEVICES = DeviceSet(["pfvm:cpu"])
49EMU_DEVICES = DeviceSet(["emu2"])
50MLSDK_DEVICES = MNCORE_DEVICES + PFVM_DEVICES + EMU_DEVICES
51
52
53def decode_loop(
54 text_decoder_embedding: TextDecoderEmbedding,
55 decoder_fn: Callable | CompiledFunction,
56 decoder_inputs: dict[str, TensorLike],
57 decoding_task: DecodingTask,
58 tokens: torch.Tensor,
59) -> tuple[torch.Tensor, torch.Tensor]:
60 sum_logprobs = torch.zeros(args.batch_size)
61 mask = torch.full(
62 (1, 1, 1, decoding_task.sample_len), -np.inf, device=tokens.device
63 )
64 max_cached_positions = mask.shape[-1]
65 current_offset = 0
66 model_dims = decoding_task.model.dims
67
68 for _ in range(decoding_task.sample_len):
69 current_seq_len = tokens.shape[1]
70 logits = None
71
72 # iterate over 2nd dim to handle variable len for the tokens tensor
73 # e.g. (batch, n) -> n x (batch, 1)
74 while current_offset < current_seq_len:
75 if current_offset >= max_cached_positions:
76 return tokens, sum_logprobs
77
78 # treat only one token in the forward pass per a iteration
79 input_tokens = tokens[:, current_offset : current_offset + 1]
80 # update mask for scaled_dot_product_attention in self_attn
81 mask[0, 0, 0, current_offset] = 0.0
82
83 # TextDecoder forward pass
84 current_offset_tensor = torch.tensor([current_offset])
85 embedded_tokens = text_decoder_embedding(
86 input_tokens, current_offset_tensor
87 )
88
89 if isinstance(decoder_inputs["embedded_tokens"], TensorProxy):
90 decoder_inputs["embedded_tokens"].load_from(
91 embedded_tokens, clone=False
92 )
93 # mask is mutated at every offset; snapshot it for asynchronous H2D.
94 decoder_inputs["mask"].load_from(mask, clone=True)
95 decoder_inputs["offset"].load_from(current_offset_tensor, clone=False)
96 else:
97 # CPU/CUDA execution calls decoder_forward directly, so provide
98 # tensors instead of MLSDK's device-resident input proxies.
99 decoder_inputs["embedded_tokens"] = embedded_tokens
100 decoder_inputs["mask"] = mask
101 decoder_inputs["offset"] = current_offset_tensor
102
103 outputs = decoder_fn(decoder_inputs)
104 logits = outputs["logits"]
105
106 current_offset += 1
107
108 # transfer logits on device to host
109 logits = logits.cpu()
110 # consider the logits at the last token only
111 logits = logits[:, -1]
112 # apply logit filters
113 for logit_filter in decoding_task.logit_filters:
114 logit_filter.apply(logits, tokens)
115
116 # expand the tokens tensor with the selected next tokens
117 # and judge if all sequences has reached the EOT.
118 tokens, completed = decoding_task.decoder.update(tokens, logits, sum_logprobs)
119
120 if completed or tokens.shape[-1] > model_dims.n_text_ctx:
121 break
122
123 return tokens, sum_logprobs
124
125
126def decode_hypotheses_texts(
127 tokens: torch.Tensor,
128 sum_logprobs: torch.Tensor,
129 decoding_task: DecodingTask,
130) -> list[str]:
131 tokens = tokens.reshape(args.batch_size, decoding_task.n_group, -1)
132 sum_logprobs = sum_logprobs.reshape(args.batch_size, decoding_task.n_group)
133
134 # get the final candidates for each group, and slice between the first sampled token and EOT
135 tokens, sum_logprobs = decoding_task.decoder.finalize(tokens, sum_logprobs)
136
137 def slice_until_eot(t: torch.Tensor) -> torch.Tensor:
138 eot_positions = (t == decoding_task.tokenizer.eot).nonzero()
139 end = int(eot_positions[0, 0]) if eot_positions.numel() else t.shape[0]
140 return t[decoding_task.sample_begin : end]
141
142 tokens = [[slice_until_eot(t) for t in s] for s in tokens]
143
144 # select the top-ranked sample in each group
145 selected = decoding_task.sequence_ranker.rank(tokens, sum_logprobs)
146 tokens = [t[i].tolist() for i, t in zip(selected, tokens)]
147 texts = [decoding_task.tokenizer.decode(t).strip() for t in tokens]
148
149 return texts
150
151
152def transcribe_mels( # noqa: CFQ002
153 args: argparse.Namespace,
154 audio_encoder_head: AudioEncoderHead,
155 encoder_fn: Callable | CompiledFunction,
156 text_decoder_embedding: TextDecoderEmbedding,
157 text_decoder_transformer: TextDecoderTransformer,
158 decoder_fn: Callable | CompiledFunction,
159 decoder_inputs: dict[str, TensorLike],
160 dataloader: torch.utils.data.DataLoader,
161 encoder_sample: dict[str, torch.Tensor],
162 decoding_task: DecodingTask,
163 context: Context | None,
164) -> tuple[float, list[float]]:
165
166 if context is not None:
167 assert isinstance(decoder_inputs["embedded_tokens"], TensorProxy)
168
169 hypotheses = []
170 references = []
171
172 with Timer() as t:
173
174 total_iterations = len(dataloader)
175 if args.max_iterations != -1:
176 total_iterations = min(total_iterations, args.max_iterations)
177
178 for mels, ref_texts in tqdm(
179 islice(dataloader, total_iterations), total=total_iterations
180 ):
181
182 if mels.ndim == 2:
183 mels = mels.unsqueeze(0)
184
185 decoding_task.decoder.reset()
186 text_decoder_transformer.reset_self_kv_cache(context)
187
188 # whisper.encoder.forward pass
189 embeddings = audio_encoder_head(mels)
190 encoder_sample.update(embeddings=embeddings)
191 cross_kv_cache = encoder_fn(encoder_sample)
192
193 # Compiled encoder outputs are TensorProxy objects and can be passed
194 # directly to the compiled decoder without a device-to-host-to-device
195 # round trip. Native torch backends use the same mapping with tensors.
196 decoder_inputs.update(cross_kv_cache)
197
198 tokens = torch.tensor([decoding_task.initial_tokens]).repeat(
199 args.batch_size, 1
200 )
201 tokens = tokens.repeat_interleave(decoding_task.n_group, dim=0)
202
203 # decode tokens and cumulate
204 tokens, sum_logprobs = decode_loop(
205 text_decoder_embedding,
206 decoder_fn,
207 decoder_inputs,
208 decoding_task,
209 tokens,
210 )
211
212 # decode texts from the inferred results
213 texts = decode_hypotheses_texts(tokens, sum_logprobs, decoding_task)
214 hypotheses.extend(texts)
215 references.extend(ref_texts)
216
217 # Calc Word Error Rate (WER) from the inferred result and the reference data
218 wer = calc_wer(hypotheses, references)
219
220 return wer, [t.time]
221
222
223def run_infer( # noqa: CFQ001, CFQ002
224 args: argparse.Namespace,
225 audio_encoder_head: AudioEncoderHead,
226 audio_encoder_transformer: AudioEncoderTransformer,
227 text_decoder_embedding: TextDecoderEmbedding,
228 text_decoder_transformer: TextDecoderTransformer,
229 decoding_task: DecodingTask,
230 dataloader: torch.utils.data.DataLoader,
231 context: Context | None,
232 outdir: str,
233) -> tuple[float, list[float]]:
234
235 model_dims = decoding_task.model.dims
236 cross_kv_precompute = CrossKVPrecompute(text_decoder_transformer)
237
238 def encoder_forward(sample_d: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
239 with torch.no_grad():
240 audio_features = audio_encoder_transformer(sample_d["embeddings"])
241 cross_kv_cache = cross_kv_precompute(audio_features)
242 return cross_kv_cache
243
244 def decoder_forward(sample_d: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
245 cross_kv_cache = kv_cache_from_dict(sample_d, model_dims.n_text_layer, "cross")
246
247 with torch.no_grad():
248 logits = text_decoder_transformer(
249 sample_d["embedded_tokens"],
250 sample_d["offset"],
251 sample_d["mask"],
252 cross_kv_cache,
253 )
254
255 return {"logits": logits}
256
257 # prepare sample inputs for compilation of encoder_fn and decoder_fn
258 # The decoder processes one token per invocation in decode_loop().
259 decode_token_length = 1
260 encoder_sample = {
261 "embeddings": torch.randn(
262 args.batch_size, model_dims.n_audio_ctx, model_dims.n_audio_state
263 )
264 }
265 cross_k_cache_dict, cross_v_cache_dict = prepare_cache_sample(args, model_dims)
266
267 decoder_sample = (
268 {
269 "embedded_tokens": torch.randn(
270 args.batch_size,
271 decode_token_length,
272 model_dims.n_text_state,
273 dtype=torch.float32,
274 ),
275 "offset": torch.tensor([1]),
276 "mask": torch.zeros(
277 1, 1, 1, model_dims.n_text_ctx // 2
278 ), # size of dim=3 is hard-coded
279 }
280 | cross_k_cache_dict
281 | cross_v_cache_dict
282 )
283
284 encoder_fn = encoder_forward
285 decoder_fn = decoder_forward
286 decoder_inputs: dict[str, TensorLike] = dict(decoder_sample)
287 if context is not None: # for mlsdk
288 # CrossKVPrecompute shares key/value projection parameters with the
289 # decoder. Register the owning decoder module first to avoid registering
290 # the shared tensors twice while compiling the encoder.
291 register_model(context, "whisper_decoder_transformer", text_decoder_transformer)
292
293 # Compile the consumer first so its preferred cross-attention cache
294 # layouts become available to the encoder compilation.
295 print("\n########## compile decoder_fn ###########")
296 decoder_fn = compile_fn(
297 context,
298 decoder_fn,
299 text_decoder_transformer,
300 decoder_sample,
301 outdir=outdir + "_decoder_transformer",
302 model_name="whisper_decoder_transformer",
303 is_train=False,
304 optimize_option=args.optimize_option,
305 preset_options_dir=args.preset_options_dir,
306 enable_cache=not args.disable_compile_cache,
307 )
308
309 decoder_inputs = decoder_fn.allocate_input_proxy()
310 cross_cache_names = set(cross_k_cache_dict) | set(cross_v_cache_dict)
311 for name, tensor in decoder_sample.items():
312 if name not in cross_cache_names:
313 decoder_inputs[name].load_from(tensor, clone=False)
314
315 # Constrain the encoder outputs to the decoder input IOSpecs. This passes
316 # only layout/type metadata to compilation; the encoder and decoder keep
317 # independent allocations, and runtime relocation connects them directly.
318 cross_cache_constraints = {
319 name: decoder_fn.input_specs[name] for name in cross_cache_names
320 }
321
322 print("\n########## compile encoder_fn ##########")
323 encoder_fn = compile_fn(
324 context,
325 encoder_fn,
326 {
327 "whisper_encoder_transformer": audio_encoder_transformer,
328 "whisper_cross_kv_precompute": cross_kv_precompute,
329 },
330 encoder_sample,
331 outdir=outdir + "_encoder_transformer",
332 model_name="whisper_encoder_transformer",
333 is_train=False,
334 optimize_option=args.optimize_option,
335 preset_options_dir=args.preset_options_dir,
336 enable_cache=not args.disable_compile_cache,
337 io_spec_constraints=cross_cache_constraints,
338 )
339
340 return transcribe_mels(
341 args,
342 audio_encoder_head,
343 encoder_fn,
344 text_decoder_embedding,
345 text_decoder_transformer,
346 decoder_fn,
347 decoder_inputs,
348 dataloader,
349 encoder_sample,
350 decoding_task,
351 context,
352 )
353
354
355def main(args: argparse.Namespace) -> None:
356 if args.max_iterations == 0 or args.max_iterations < -1:
357 raise ValueError("max_iterations must be a positive integer or -1")
358 if args.wer_criteria < 0.0 or args.wer_criteria > 1.0:
359 raise ValueError("wer_criteria must be ranged in [0, 1]")
360
361 # Fix seed values for reproducibility
362 set_deterministic_mode(seed=args.seed)
363
364 # Create model obj
365 model = whisper.load_model(args.model, download_root=args.model_dir)
366 model.alignment_heads = model.alignment_heads.to_dense()
367 model.eval()
368 audio_encoder_head = AudioEncoderHead(model.encoder)
369 audio_encoder_transformer = AudioEncoderTransformer(model.encoder)
370 text_decoder_embedding = TextDecoderEmbedding(model.decoder)
371 text_decoder_transformer = TextDecoderTransformer(model.decoder, args.batch_size)
372
373 # Create decoding_task obj to use DecodingTask params in the inference
374 decoding_options = DecodingOptions(language="en", without_timestamps=True)
375 decoding_task = DecodingTask(model, decoding_options)
376
377 # Create Dataset/DataLoader obj
378 infer_dataset = LibriSpeech(
379 dataset_dir=args.dataset_dir,
380 data_name=args.data_name,
381 chunk_length=args.chunk_length,
382 n_mels=model.dims.n_mels,
383 device="cpu",
384 )
385 infer_dataloader = torch.utils.data.DataLoader(
386 infer_dataset, batch_size=args.batch_size, drop_last=True
387 )
388 infer_iterations = len(infer_dataloader)
389 if args.max_iterations != -1:
390 infer_iterations = min(infer_iterations, args.max_iterations)
391
392 # Decide device and outdir from given command line args/options
393 sample_name = "whisper_inference"
394 outdir = decide_outdir(
395 args.backend, example_name=sample_name, basedir=args.out_basedir
396 )
397
398 # Pass device info to the Context obj
399 context = None
400 if args.backend in MLSDK_DEVICES:
401 device = MNDevice(args.backend)
402 context = Context(device)
403 Context.switch_context(context)
404 else:
405 device = args.backend
406
407 with trace_scope(args.trace_filename):
408 # Run inference
409 wer, infer_times = run_infer(
410 args,
411 audio_encoder_head,
412 audio_encoder_transformer,
413 text_decoder_embedding,
414 text_decoder_transformer,
415 decoding_task,
416 infer_dataloader,
417 context,
418 outdir,
419 )
420
421 # Output the inference time
422 output_result_times(
423 eval_times=infer_times,
424 eval_iter=infer_iterations,
425 eval_batch_size=args.batch_size,
426 backend_name=args.backend,
427 sample_name=sample_name,
428 optimize_option=args.optimize_option,
429 )
430
431 print(f"\nResult: WordErrorRate(WER) = {wer * 100:.2f} %")
432 assert wer <= args.wer_criteria, (
433 f"{wer * 100:.2f} % exceeded the "
434 f"specified criteria ({args.wer_criteria * 100:.2f} %)."
435 )
436
437
438if __name__ == "__main__":
439 script_dir = os.path.dirname(__file__)
440 parser = argparse.ArgumentParser(
441 formatter_class=argparse.ArgumentDefaultsHelpFormatter
442 )
443
444 parser.add_argument(
445 "--dataset_dir",
446 type=str,
447 default=f"{script_dir}/librispeech_dataset",
448 help="path to save dataset",
449 )
450 parser.add_argument(
451 "--out_basedir",
452 type=str,
453 default="/tmp",
454 help="basedir to save output artifacts",
455 )
456 parser.add_argument(
457 "--model_dir",
458 type=str,
459 default="/tmp/whisper_models",
460 help="directory in which to download and load Whisper models",
461 )
462 parser.add_argument(
463 "--max_iterations",
464 type=int,
465 default=-1,
466 help="maximum number of DataLoader iterations; -1 processes all samples",
467 )
468 parser.add_argument(
469 "--wer_criteria",
470 type=float,
471 default=0.05,
472 help="acceptable WordErrorRate (WER) criteria",
473 )
474
475 # mlsdk options
476 parser.add_argument(
477 "-b",
478 "--backend",
479 type=str,
480 choices=DEVICES,
481 default="cpu",
482 help="mlsdk options: specify a device as the backend (default: %(default)s)",
483 )
484 parser.add_argument(
485 "--optimize_option",
486 type=str,
487 default="O1",
488 choices=["debug", "O0", "O1", "O2", "O3", "O4"],
489 help="mlsdk options: optimize option for Context.compile() (default %(default)s)",
490 )
491 parser.add_argument(
492 "--preset_options_dir",
493 type=str,
494 default="/opt/pfn/pfcomp/codegen/preset_options",
495 help="mlsdk options: path to preset_options/",
496 )
497 parser.add_argument(
498 "--disable_compile_cache",
499 action="store_true",
500 help="mlsdk options: disable reusing compiled artifacts",
501 )
502 parser.add_argument(
503 "--trace_filename",
504 type=str,
505 default=None,
506 help="mlsdk options: specify a valid filename to enable tracing",
507 )
508
509 apply_toml_defaults(f"{script_dir}/configs.toml", parser)
510
511 args = parser.parse_args()
512
513 main(args)
1# Portions adapted from openai/whisper. Copyright (c) 2022 OpenAI.
2# See LICENSE.whisper for the MIT License and license_note.md for source references.
3
4import argparse
5import os
6
7import jiwer
8import torch
9import torchaudio
10import whisper
11from whisper.model import ModelDimensions
12from whisper.normalizers import EnglishTextNormalizer
13
14
15# Wrapper class for torchaudio.datasets.LIBRISPEECH class
16# class definition is originally from
17# [whisper example](https://github.com/openai/whisper/blob/c0d2f624c09dc18e709e37c2ad90c039a4eb72a2/notebooks/LibriSpeech.ipynb#L56) # noqa: B950
18class LibriSpeech(torch.utils.data.Dataset):
19 def __init__(
20 self,
21 dataset_dir: str,
22 data_name: str = "test-clean",
23 chunk_length: int = 30,
24 n_mels: int = 80,
25 device: str | torch.device = "cpu",
26 ) -> None:
27 if not os.path.isdir(dataset_dir):
28 os.makedirs(dataset_dir)
29 self.dataset = torchaudio.datasets.LIBRISPEECH(
30 root=dataset_dir,
31 url=data_name,
32 download=True,
33 )
34 self.device = device
35 self.chunk_length = chunk_length
36 self.n_mels = n_mels
37 self.SAMPLE_RATE = 16000 # Fixed param
38
39 def __len__(self) -> int:
40 return len(self.dataset)
41
42 def __getitem__(self, item: int) -> tuple[torch.Tensor, str]:
43 audio, sample_rate, text, _, _, _ = self.dataset[item]
44 assert sample_rate == self.SAMPLE_RATE
45 audio = whisper.pad_or_trim(
46 audio.flatten(), length=self.SAMPLE_RATE * self.chunk_length
47 ).to(self.device)
48 mel = whisper.log_mel_spectrogram(audio, n_mels=self.n_mels)
49
50 return (mel, text)
51
52
53def prepare_cache_sample(
54 args: argparse.Namespace,
55 model_dims: ModelDimensions,
56 dtype: torch.dtype = torch.float32,
57) -> tuple[
58 dict[str, torch.Tensor],
59 dict[str, torch.Tensor],
60]:
61
62 def create_cache_sample(
63 size: list[int],
64 n_text_layer: int,
65 cache_prefix: str,
66 dtype: torch.dtype = torch.float32,
67 ) -> dict[str, torch.Tensor]:
68 return {
69 f"{cache_prefix}_cache_{i}": torch.zeros(*size, dtype=dtype)
70 for i in range(n_text_layer)
71 }
72
73 # cache sizes
74 n_head = model_dims.n_text_head
75 n_state_per_head = model_dims.n_text_state // model_dims.n_text_head
76 n_layer = model_dims.n_text_layer
77 cross_kv_cache_size = [
78 args.batch_size,
79 n_head,
80 model_dims.n_audio_ctx,
81 n_state_per_head,
82 ]
83
84 cross_k_cache_dict = create_cache_sample(
85 cross_kv_cache_size, n_layer, "cross_k", dtype
86 )
87 cross_v_cache_dict = create_cache_sample(
88 cross_kv_cache_size, n_layer, "cross_v", dtype
89 )
90
91 return cross_k_cache_dict, cross_v_cache_dict
92
93
94def kv_cache_from_dict(
95 sample_d: dict[str, torch.Tensor],
96 n_text_layer: int,
97 cache_prefix: str, # self or cross
98) -> list[tuple[torch.Tensor, torch.Tensor]]:
99
100 kv_cache = [
101 (
102 sample_d[f"{cache_prefix}_k_cache_{i}"],
103 sample_d[f"{cache_prefix}_v_cache_{i}"],
104 )
105 for i in range(n_text_layer)
106 ]
107
108 return kv_cache
109
110
111def calc_wer(
112 hypotheses: list[str], references: list[str], print_texts: bool = True
113) -> float:
114 normalizer = EnglishTextNormalizer()
115 hypotheses = [normalizer(text) for text in hypotheses]
116 references = [normalizer(text) for text in references]
117
118 if print_texts:
119 print("\n############ Outputs #############")
120 print("\nNormalized references and hypotheses (first 3 samples):")
121 for index, (reference, hypothesis) in enumerate(
122 zip(references[:3], hypotheses[:3])
123 ):
124 print(
125 f"[{index}]\n"
126 f" Reference : {reference}\n"
127 f" Hypothesis: {hypothesis}"
128 )
129 print("\n##################################")
130
131 return jiwer.wer(references, hypotheses)
1# Portions adapted from openai/whisper. Copyright (c) 2022 OpenAI.
2# See LICENSE.whisper for the MIT License and license_note.md for source references.
3
4import types
5
6import torch
7from whisper.model import AudioEncoder, TextDecoder
8
9
10# to whisper.encoder's head modules (conv1(Conv1d) and conv2(Conv1d))
11class AudioEncoderHead(torch.nn.Module):
12 def __init__(self, org_encoder: AudioEncoder) -> None:
13 super().__init__()
14
15 self.conv1 = org_encoder.conv1
16 self.conv2 = org_encoder.conv2
17
18 def forward(self, x: torch.Tensor) -> torch.Tensor:
19 x = torch.nn.functional.gelu(self.conv1(x))
20 x = torch.nn.functional.gelu(self.conv2(x))
21 x = x.permute(0, 2, 1)
22
23 return x
24
25
26def encoder_self_attn_forward(self, x: torch.Tensor) -> torch.Tensor:
27 q = self.query(x)
28 k = self.key(x)
29 v = self.value(x)
30
31 q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
32 k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
33 v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
34
35 a = torch.nn.functional.scaled_dot_product_attention(q, k, v)
36 out = a.permute(0, 2, 1, 3).flatten(start_dim=2)
37
38 return self.out(out)
39
40
41def encoder_block_forward(self, x: torch.Tensor) -> torch.Tensor:
42 x = x + self.attn(self.attn_ln(x))
43 x = x + self.mlp(self.mlp_ln(x))
44 return x
45
46
47# whisper.encoder's module except to the head
48class AudioEncoderTransformer(torch.nn.Module):
49 def __init__(self, org_encoder: AudioEncoder) -> None:
50 super().__init__()
51
52 self.register_buffer("positional_embedding", org_encoder.positional_embedding)
53 self.blocks = org_encoder.blocks
54 self.ln_post = org_encoder.ln_post
55
56 for block in self.blocks:
57 # replace ResidualAttentionBlock.forward()
58 block.forward = types.MethodType(encoder_block_forward, block)
59 # replace MultiHeadAttention.forward()
60 block.attn.forward = types.MethodType(encoder_self_attn_forward, block.attn)
61
62 def forward(self, x: torch.Tensor) -> torch.Tensor:
63
64 assert (
65 x.shape[1:] == self.positional_embedding.shape
66 ), f"incorrect audio shape: {x.shape=} vs. {self.positional_embedding.shape=}"
67
68 x = (x + self.positional_embedding).to(x.dtype)
69
70 for block in self.blocks:
71 x = block(x)
72
73 x = self.ln_post(x)
74 return x
75
76
77def static_self_attn_forward(
78 self,
79 x: torch.Tensor,
80 mask: torch.Tensor,
81 kv_cache: tuple[torch.Tensor, torch.Tensor],
82 offset: torch.Tensor,
83) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
84 q = self.query(x)
85 k = self.key(x)
86 v = self.value(x)
87
88 q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
89 k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
90 v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
91
92 k_cache = kv_cache[0]
93 v_cache = kv_cache[1]
94
95 # Update the persistent KV cache in place. The cache tensors are registered
96 # buffers of TextDecoderTransformer, so MLSDK keeps their device allocation
97 # across decoder invocations.
98 offset_idx = offset.view(1, 1, 1, 1).expand(*q.shape)
99 k_cache.scatter_(dim=2, index=offset_idx, src=k)
100 v_cache.scatter_(dim=2, index=offset_idx, src=v)
101
102 a = torch.nn.functional.scaled_dot_product_attention(
103 q,
104 k_cache,
105 v_cache,
106 attn_mask=mask,
107 dropout_p=0.0,
108 is_causal=False,
109 )
110 out = a.permute(0, 2, 1, 3).flatten(start_dim=2)
111
112 return self.out(out)
113
114
115def static_cross_attn_forward(
116 self,
117 x: torch.Tensor,
118 kv_cache: tuple[torch.Tensor, torch.Tensor],
119) -> torch.Tensor:
120 q = self.query(x)
121 q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
122 # view() and permute() have been done in CrossKVPrecompute.
123 k, v = kv_cache
124
125 a = torch.nn.functional.scaled_dot_product_attention(q, k, v)
126 out = a.permute(0, 2, 1, 3).flatten(start_dim=2)
127
128 return self.out(out)
129
130
131def static_block_forward(
132 self,
133 x: torch.Tensor,
134 mask: torch.Tensor | None = None,
135 kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
136 offset: torch.Tensor | None = None,
137 cross_kv_cache: (
138 tuple[torch.Tensor, torch.Tensor] | None
139 ) = None, # should be precomputed
140) -> torch.Tensor:
141 # apply self_attention
142 attn_out = self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache, offset=offset)
143 x = x + attn_out
144
145 # apply cross_attention
146 if self.cross_attn:
147 x = x + self.cross_attn(self.cross_attn_ln(x), kv_cache=cross_kv_cache)
148
149 # apply mlp
150 x = x + self.mlp(self.mlp_ln(x))
151
152 return x
153
154
155# to calculation of the embeddings of the input tokens
156class TextDecoderEmbedding(torch.nn.Module):
157 def __init__(self, org_decoder: TextDecoder) -> None:
158 super().__init__()
159 self.token_embedding = org_decoder.token_embedding
160 self.positional_embedding = org_decoder.positional_embedding
161
162 def forward(self, x: torch.Tensor, offset: torch.Tensor) -> torch.Tensor:
163 pos_emb = self.positional_embedding[offset[0]].view(
164 1, 1, -1
165 ) # to broadcast explictly
166 x = self.token_embedding(x) + pos_emb
167
168 return x
169
170
171# whisper.decoder's module except to embeddings
172class TextDecoderTransformer(torch.nn.Module):
173 def __init__(self, org_decoder: TextDecoder, batch_size: int) -> None:
174 super().__init__()
175
176 self.blocks = org_decoder.blocks
177 self.ln = org_decoder.ln
178 self.token_embedding = org_decoder.token_embedding
179
180 # TextDecoder does not expose n_ctx; its positional embedding has one
181 # row per text-context position.
182 cache_length = org_decoder.positional_embedding.shape[0] // 2
183 n_head = self.blocks[0].attn.n_head
184 n_state_per_head = self.token_embedding.embedding_dim // n_head
185 cache_shape = (batch_size, n_head, cache_length, n_state_per_head)
186 for i in range(len(self.blocks)):
187 self.register_buffer(f"self_k_cache_{i}", torch.zeros(cache_shape))
188 self.register_buffer(f"self_v_cache_{i}", torch.zeros(cache_shape))
189
190 for block in self.blocks:
191 # replace forward() of ResidualAttentionBlock
192 block.forward = types.MethodType(static_block_forward, block)
193
194 # replace forward() of self_attention
195 block.attn.forward = types.MethodType(static_self_attn_forward, block.attn)
196
197 # replace forward() of cross_attention
198 block.cross_attn.forward = types.MethodType(
199 static_cross_attn_forward, block.cross_attn
200 )
201
202 def self_kv_cache(self) -> list[tuple[torch.Tensor, torch.Tensor]]:
203 return [
204 (
205 getattr(self, f"self_k_cache_{i}"),
206 getattr(self, f"self_v_cache_{i}"),
207 )
208 for i in range(len(self.blocks))
209 ]
210
211 def reset_self_kv_cache(self, context=None) -> None:
212 """Clear self-attention state before decoding the next audio sample."""
213 for k_cache, v_cache in self.self_kv_cache():
214 k_cache.zero_()
215 v_cache.zero_()
216 if context is not None:
217 context.get_registered_value_proxy(k_cache).load_from(
218 k_cache, clone=False
219 )
220 context.get_registered_value_proxy(v_cache).load_from(
221 v_cache, clone=False
222 )
223
224 def forward(
225 self,
226 x: torch.Tensor, # embedded_tokens
227 offset: torch.Tensor, # offset for kv_cache
228 mask: torch.Tensor, # mask for kv_cache
229 cross_kv_cache: list[tuple[torch.Tensor, torch.Tensor]], # len(list) == n_layer
230 ) -> torch.Tensor:
231
232 x = x.to(cross_kv_cache[0][0].dtype)
233
234 self_kv_cache = self.self_kv_cache()
235 for i, block in enumerate(self.blocks):
236 self_kv = self_kv_cache[i]
237 cross_kv = cross_kv_cache[i]
238 x = block(
239 x,
240 mask=mask,
241 kv_cache=self_kv,
242 offset=offset,
243 cross_kv_cache=cross_kv,
244 )
245
246 x = self.ln(x)
247 logits = (
248 x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)
249 ).float()
250
251 return logits
252
253
254class CrossKVPrecompute(torch.nn.Module):
255 """Precompute decoder cross-attention K/V caches from audio features."""
256
257 def __init__(self, text_decoder_transformer: TextDecoderTransformer) -> None:
258 super().__init__()
259
260 # These projection modules are shared with TextDecoderTransformer.
261 # Register TextDecoderTransformer with the MLSDK Context before compiling
262 # the encoder so that their parameters are registered only once.
263 self.keys = torch.nn.ModuleList(
264 [block.cross_attn.key for block in text_decoder_transformer.blocks]
265 )
266 self.values = torch.nn.ModuleList(
267 [block.cross_attn.value for block in text_decoder_transformer.blocks]
268 )
269 self.n_heads = [
270 block.cross_attn.n_head for block in text_decoder_transformer.blocks
271 ]
272
273 def forward(self, audio_features: torch.Tensor) -> dict[str, torch.Tensor]:
274 cross_kv_cache = {}
275
276 for i, (key, value, n_head) in enumerate(
277 zip(self.keys, self.values, self.n_heads)
278 ):
279 k = key(audio_features)
280 v = value(audio_features)
281
282 k = k.view(k.shape[0], k.shape[1], n_head, -1).permute(0, 2, 1, 3)
283 v = v.view(v.shape[0], v.shape[1], n_head, -1).permute(0, 2, 1, 3)
284
285 cross_kv_cache[f"cross_k_cache_{i}"] = k
286 cross_kv_cache[f"cross_v_cache_{i}"] = v
287
288 return cross_kv_cache
1import argparse
2import gc
3import inspect
4import math
5import os
6import random
7import re
8import sys
9import time
10from collections.abc import Callable
11from typing import Any, Self
12
13import numpy as np
14import tomllib
15import torch
16
17try:
18 from fx2onnx import set_tensor_name
19 from mlsdk import (
20 CacheOptions,
21 CompiledFunction,
22 Context,
23 MNCoreOptimizer,
24 TensorProxy,
25 get_tensor_name,
26 set_buffer_name_in_optimizer,
27 set_tensor_name_in_module,
28 storage,
29 )
30except Exception as e:
31 # For image without mlsdk
32 print(f"Warning: {e}")
33
34 # dummies for missing dependency
35 class CompiledFunction:
36 pass
37
38 class Context:
39 pass
40
41 class MNCoreOptimizer:
42 pass
43
44 class TensorProxy:
45 pass
46
47 pass
48
49
50class DeviceSet:
51 def __init__(self, patterns):
52 self.patterns = patterns
53
54 def __contains__(self, item):
55 return any(bool(re.match(pattern, item)) for pattern in self.patterns)
56
57 def __iter__(self):
58 for pattern in self.patterns:
59 yield pattern
60
61 def __add__(self, other):
62 return DeviceSet(self.patterns + other.patterns)
63
64
65def decide_outdir(
66 device: str,
67 example_name: str = "example",
68 basedir: str | os.PathLike = ".",
69) -> str:
70 outdir = f"{basedir}/outdir_{example_name}_"
71 run_without_mlsdk = False
72 if "mncore2" in device:
73 outdir += "mncore"
74 elif "pfvm:cpu" in device:
75 outdir += "pfvm_cpu"
76 elif "pfvm:cuda" in device:
77 outdir += "pfvm_cuda"
78 elif "emu" in device:
79 outdir += "mncore_emu2"
80 elif "cuda" in device:
81 outdir += "cuda"
82 run_without_mlsdk = True
83 else:
84 outdir += "cpu"
85 run_without_mlsdk = True
86
87 # Create output dirs when using "cpu" and "cuda" backends to save the outputs
88 if run_without_mlsdk:
89 os.makedirs(outdir, exist_ok=True)
90
91 return outdir
92
93
94def register_model(
95 context: Context,
96 name: str,
97 model: torch.nn.Module,
98) -> None:
99 if (
100 get_tensor_name(next(model.parameters())) is None
101 ): # in case the model obj isn't registered to the context
102 set_tensor_name_in_module(model, name)
103 for p in model.parameters():
104 context.register_param(p)
105 for b in model.buffers():
106 context.register_buffer(b)
107
108
109def compile_fn( # noqa: CFQ002
110 context: Context,
111 target_fn: Callable[
112 [
113 dict[str, torch.Tensor],
114 ],
115 dict[str, torch.Tensor],
116 ], # compiled fn
117 model: torch.nn.Module | dict[str, torch.nn.Module],
118 sample_input: dict[str, torch.Tensor],
119 outdir: str = "/tmp/example_output",
120 model_name: str = "example",
121 is_train: bool = True,
122 optimizers: (
123 list[MNCoreOptimizer | torch.optim.Optimizer] | None
124 ) = None, # list[] is for multiple optimizers
125 optimize_option: str = "debug",
126 preset_options_dir: str = "/opt/pfn/pfcomp/codegen/preset_options",
127 enable_cache: bool = False,
128 **kwargs: Any, # used in `compile_args` in Context.compile()
129) -> CompiledFunction:
130
131 compile_options = {
132 "option_json": os.path.join(preset_options_dir, optimize_option + ".json")
133 }
134
135 compile_args = {
136 "function": target_fn,
137 "inputs": sample_input,
138 "options": compile_options,
139 }
140
141 codegen_base_dir = storage.path(outdir)
142 if needs_old_compile_args():
143 compile_args["codegen_base_dir"] = codegen_base_dir
144 compile_args["name"] = model_name
145 else: # for newer MLSDK
146 compile_args["codegen_dir"] = codegen_base_dir / model_name
147
148 if enable_cache:
149 compile_args["cache_options"] = CacheOptions(
150 f"{outdir}/cache",
151 enable_app_cache=True,
152 enable_onnx_cache=True,
153 enable_codegen_cache=True,
154 enable_gpfn2obj_cache=True,
155 )
156
157 if isinstance(model, torch.nn.Module):
158 register_model(context, model_name, model)
159 else: # if isinstance(models, dict[str, torch.nn.Module]):
160 for name, actual_model in model.items():
161 register_model(context, name, actual_model)
162
163 if is_train:
164 if optimizers is None: # in case that optimizer.step() will be done at the host
165 if isinstance(model, torch.nn.Module):
166 for n, p in model.named_parameters():
167 p.grad = torch.nn.Parameter(
168 torch.zeros_like(p), requires_grad=p.requires_grad
169 )
170 set_tensor_name(p.grad, f"{model_name}@{n}@grad".replace(".", "_"))
171 context.register_param(p.grad)
172 else:
173 for name, actual_model in model.items():
174 for n, p in actual_model.named_parameters():
175 p.grad = torch.nn.Parameter(
176 torch.zeros_like(p), requires_grad=p.requires_grad
177 )
178 set_tensor_name(p.grad, f"{name}@{n}".replace(".", "_"))
179 context.register_param(p.grad)
180 else:
181 for idx, optimizer in enumerate(optimizers):
182 optimizer_name = "optimizer" + str(idx)
183 set_buffer_name_in_optimizer(optimizer, optimizer_name)
184 context.register_optimizer_buffers(optimizer)
185
186 compile_args.update(kwargs)
187
188 return context.compile(**compile_args)
189
190
191class Timer:
192 def __init__(self) -> None:
193 self.time = None
194
195 def __enter__(self) -> Self:
196 gc.disable()
197 self.start_time = time.perf_counter()
198 return self
199
200 def __exit__(self, exc_type, exc_value, traceback) -> None:
201 end_time = time.perf_counter()
202 gc.enable()
203 self.time = end_time - self.start_time
204 return None
205
206
207def set_deterministic_mode(seed: int) -> None:
208 # Set seed
209 random.seed(seed)
210 np.random.seed(seed)
211 torch.manual_seed(seed)
212 torch.cuda.manual_seed(seed)
213
214 # Set cudnn.benchmark mode and specify the use of deterministic algorithms
215 torch.backends.cudnn.benchmark = False
216 torch.use_deterministic_algorithms(True)
217
218
219def output_result_times( # noqa: CFQ002
220 train_times: list[float] | None = None,
221 eval_times: list[float] | None = None,
222 train_iter: int | None = None,
223 eval_iter: int | None = None,
224 train_batch_size: int | None = None,
225 eval_batch_size: int | None = None,
226 optimizer_name: str | None = None,
227 backend_name: str = "cpu",
228 sample_name: str = "sample",
229 optimize_option: str = "debug",
230 time_scale: str = "s",
231) -> None:
232 scaling_factor = 1.0
233 if time_scale == "ms":
234 scaling_factor = 1000.0
235 elif time_scale == "us":
236 scaling_factor = 1000000.0
237 elif time_scale == "ns":
238 scaling_factor = 1000000000.0
239
240 onnx_exporter = "None"
241 optimize_flag = "None"
242 if backend_name not in ["cpu", "cuda"]:
243 onnx_exporter = "fx2onnx"
244
245 optimize_flag = optimize_option
246
247 print("\n########## Performances ##########")
248 print("\n---------- configs ----------")
249 print(f"sample name: {sample_name}")
250 print(f"backend device: {backend_name}")
251 print(f"onnx exporter: {onnx_exporter}")
252 print(f"optimizer: {optimizer_name}")
253 print(f"optimize_option: {optimize_flag}")
254 print("------------------------------")
255
256 if train_times is not None and len(train_times) != 0:
257 print("\n---------- performance of training part ----------")
258 total_time = math.fsum(train_times) * scaling_factor
259 average_time = total_time / len(train_times)
260 iter_per_sec = train_iter * len(train_times) / total_time
261 print(f"train epochs: {len(train_times)}")
262 print(f"batch size: {train_batch_size}")
263 print(f"iterations per epoch: {train_iter}")
264 print(f"total time [{time_scale}]: {total_time}")
265 print(f"average per epoch [{time_scale}]: {average_time}")
266 print(f"averaged [iter/{time_scale}]: {iter_per_sec}")
267 print("--------------------------------------------------")
268
269 if eval_times is not None and len(eval_times) != 0:
270 print("\n---------- performance of evaluation/inference part ----------")
271 total_time = math.fsum(eval_times) * scaling_factor
272 average_time = total_time / len(eval_times)
273 iter_per_sec = eval_iter * len(eval_times) / total_time
274 print(f"eval epochs: {len(eval_times)}")
275 print(f"batch size: {eval_batch_size}")
276 print(f"iterations per epoch: {eval_iter}")
277 print(f"total time [{time_scale}]: {total_time}")
278 print(f"average per epoch [{time_scale}]: {average_time}")
279 print(f"averaged [iter/{time_scale}]: {iter_per_sec}")
280 print("--------------------------------------------------------------")
281 print("\n##################################")
282
283
284def needs_old_compile_args() -> bool:
285 parameter_names = inspect.signature(Context.compile).parameters.keys()
286 return "codegen_base_dir" in parameter_names
287
288
289# for type hint of the configs from toml
290class TomlValue:
291 str | int | float | bool | list["TomlValue"] | dict[str, "TomlValue"]
292
293
294class TomlDict:
295 dict[str, TomlValue]
296
297
298def read_configs_from_toml(
299 toml_path: str,
300) -> TomlDict:
301
302 configs_dict = None
303 with open(toml_path, mode="rb") as f:
304 configs_dict = tomllib.load(f)
305
306 return configs_dict
307
308
309def str2bool(v: bool | str) -> bool:
310 if v.lower() in ("yes", "true", "on", "enable", "y", "t", "1"):
311 return True
312 elif v.lower() in ("no", "false", "off", "disable", "n", "f", "0"):
313 return False
314 elif isinstance(v, str | bool):
315 return v
316 else:
317 raise argparse.ArgumentTypeError("Str or boolean value expected")
318
319
320def apply_toml_defaults(
321 configs: TomlDict | str | os.PathLike,
322 parser: argparse.ArgumentParser,
323) -> None:
324
325 if isinstance(configs, dict):
326 for k, v in configs.items():
327 if isinstance(v, dict): # in case v is (nested) dict
328 apply_toml_defaults(v, parser)
329 else:
330 # just checking whether v is list is enough for array args
331 # because array in toml is converted to the list by tomllib.
332 args_type = None
333 if isinstance(v, list):
334 args_type = type(v[0])
335 elif isinstance(v, bool):
336 args_type = str2bool
337 else:
338 args_type = type(v)
339 parser.add_argument(
340 f"--{k}",
341 default=v,
342 type=args_type,
343 nargs="*" if isinstance(v, list) else "?",
344 )
345 elif isinstance(configs, str | os.PathLike):
346 configs_dict = read_configs_from_toml(configs)
347
348 apply_toml_defaults(configs_dict, parser)
349 else:
350 sys.exit("")
1openai-whisper==20250625
2soundfile
3jiwer
4pandas
5torchaudio==2.9.0
6torchcodec==0.9.0
1torchaudio==2.9.0
2torchcodec==0.9.0
1title = "whisper_inference"
2
3
4[data]
5data_name = "test-clean" # dataset name of LibriSpeech dataset to download
6chunk_length = 30 # size of the chunk of input audio data. here, chunk_length == the number of seconds of the audio data
7
8
9[inference]
10model = "base.en" # model names used in whisper inference
11
12
13[misc]
14seed = 0
15batch_size = 1