8.2.5. Example: Speech Recognition using Whisper

Sample program demonstrating speech recognition inference with the Whisper base.en model on the LibriSpeech ASR corpus test-clean dataset. The example reports the Word Error Rate (WER) of the recognized text.

Note

The model, pretrained weights, and some source code used in this example are based on or directly sourced from openai/whisper. They are licensed under the MIT License.

The LibriSpeech ASR corpus is licensed under CC BY 4.0.

Licenses and Source-Code References

The example installs and loads the Whisper implementation and pretrained weights distributed by openai/whisper. They are provided under the MIT License.

The following parts of this example were derived from or implemented with reference to openai/whisper sources, which are also licensed under the MIT License:

  • whisper/notebooks/LibriSpeech.ipynb: a lightweight wrapper around torchaudio.datasets.LIBRISPEECH, the inference workflow, and calculation of WER from model outputs.

  • whisper/model.py: implementations of the Whisper, AudioEncoder, TextDecoder, ResidualAttentionBlock, and MultiHeadAttention classes.

  • whisper/decoding.py: implementations of the DecodingTask and PyTorchInference classes, and the decode() function.

Preparation

The example uses the LibriSpeech dataset. On the first run, torchaudio automatically downloads and extracts the selected dataset split into ./librispeech_dataset. The default split is test-clean.

If this automatic download fails, for example with a permission error, download and extract the dataset manually. See the LibriSpeech ASR corpus page for dataset download links.

$ 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

Run the following command from the example directory:

$ cd /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference
$ ./run_whisper_inference.sh [OPTION]...

For example, the following command runs speech recognition inference on MN-Core 2:

$ ./run_whisper_inference.sh -b mncore2:auto

For a quick functional check that processes only the first 10 iterations, run:

$ ./run_whisper_inference.sh -b mncore2:auto --max_iterations 10

Command-line Options

The following options are available. Options specified in configs.toml can also be passed on the command line.

  • -h, --help: Display the available command-line options.

  • --dataset_dir DATASET_DIR: Specify the directory in which to download and store the LibriSpeech dataset. The default is ./librispeech_dataset in the example directory.

  • --data_name DATA_NAME: Specify the LibriSpeech dataset split to use. The default is test-clean.

  • --chunk_length CHUNK_LENGTH: Specify the length, in seconds, of each input audio chunk. The default is 30.

  • --model MODEL: Specify the Whisper model to load. The default is base.en.

  • --model_dir MODEL_DIR: Specify the directory in which to download and load Whisper model files. The default is /tmp/whisper_models.

  • --seed SEED: Specify the random seed used to make the run reproducible. The default is 0.

  • --batch_size BATCH_SIZE: Specify the number of audio samples processed in each batch. The default is 1.

  • --max_iterations MAX_ITERATIONS: Specify the maximum number of DataLoader iterations. The default value, -1, processes the entire dataset.

  • --wer_criteria WER_CRITERIA: Specify the maximum acceptable Word Error Rate (WER). The default is 0.05 (5%). The run fails when the measured WER exceeds this value.

  • -b BACKEND, --backend BACKEND: Specify the backend used for inference. Supported values are cpu, pfvm:cpu, mncore2:auto, mncore2:[0-7], and emu2. The default is cpu.

  • --out_basedir OUT_BASEDIR: Specify the base directory for generated output artifacts. The default is /tmp.

  • --optimize_option {debug,O0,O1,O2,O3,O4}: Specify the optimization option passed to Context.compile(). The default is O1.

  • --preset_options_dir PRESET_OPTIONS_DIR: Specify the directory that contains the preset optimization-option JSON files. The default is /opt/pfn/pfcomp/codegen/preset_options; the wrapper script overrides it with the preset_options directory in the Codegen source tree.

  • --disable_compile_cache: Disable reuse of compiled artifacts. By default, compiled artifacts are reused when the backend uses MLSDK.

  • --trace_filename TRACE_FILENAME: Specify an output filename to enable tracing. Tracing is disabled by default.

Expected Output

The program prints some cases of the recognized outputs, inference performance information and the WER calculated from the recognized text and LibriSpeech reference transcripts. For example:

############ 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 %

By default, the run fails if the WER exceeds 5%; this threshold can be changed with --wer_criteria.

Scripts

Listing 8.29 /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference/run_whisper_inference.sh
 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
Listing 8.30 /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference/whisper_inference.py
  1import argparse
  2import os
  3from collections.abc import Callable
  4from itertools import islice
  5
  6import numpy as np
  7import torch
  8import whisper
  9from mlsdk import (
 10    CompiledFunction,
 11    Context,
 12    MNDevice,
 13    TensorLike,
 14    TensorProxy,
 15    trace_scope,
 16)
 17from static_whisper import (
 18    AudioEncoderHead,
 19    AudioEncoderTransformer,
 20    CrossKVPrecompute,
 21    TextDecoderEmbedding,
 22    TextDecoderTransformer,
 23)
 24from tqdm import tqdm
 25from utility import (
 26    DeviceSet,
 27    Timer,
 28    apply_toml_defaults,
 29    compile_fn,
 30    decide_outdir,
 31    output_result_times,
 32    register_model,
 33    set_deterministic_mode,
 34)
 35from whisper.decoding import DecodingOptions, DecodingTask
 36from whisper_utils import (
 37    LibriSpeech,
 38    calc_wer,
 39    kv_cache_from_dict,
 40    prepare_cache_sample,
 41)
 42
 43DEVICES = DeviceSet(["pfvm:cpu", "mncore2:auto", "mncore2:[0-7]", "emu2", "cpu"])
 44MNCORE_DEVICES = DeviceSet(["mncore2:auto", "mncore2:[0-7]"])
 45PFVM_DEVICES = DeviceSet(["pfvm:cpu"])
 46EMU_DEVICES = DeviceSet(["emu2"])
 47MLSDK_DEVICES = MNCORE_DEVICES + PFVM_DEVICES + EMU_DEVICES
 48
 49
 50def decode_loop(
 51    text_decoder_embedding: TextDecoderEmbedding,
 52    decoder_fn: Callable | CompiledFunction,
 53    decoder_inputs: dict[str, TensorLike],
 54    decoding_task: DecodingTask,
 55    tokens: torch.Tensor,
 56) -> tuple[torch.Tensor, torch.Tensor]:
 57    sum_logprobs = torch.zeros(args.batch_size)
 58    mask = torch.full(
 59        (1, 1, 1, decoding_task.sample_len), -np.inf, device=tokens.device
 60    )
 61    max_cached_positions = mask.shape[-1]
 62    current_offset = 0
 63    model_dims = decoding_task.model.dims
 64
 65    for _ in range(decoding_task.sample_len):
 66        current_seq_len = tokens.shape[1]
 67        logits = None
 68
 69        # iterate over 2nd dim to handle variable len for the tokens tensor
 70        # e.g. (batch, n) -> n x (batch, 1)
 71        while current_offset < current_seq_len:
 72            if current_offset >= max_cached_positions:
 73                return tokens, sum_logprobs
 74
 75            # treat only one token in the forward pass per a iteration
 76            input_tokens = tokens[:, current_offset : current_offset + 1]
 77            # update mask for scaled_dot_product_attention in self_attn
 78            mask[0, 0, 0, current_offset] = 0.0
 79
 80            # TextDecoder forward pass
 81            current_offset_tensor = torch.tensor([current_offset])
 82            embedded_tokens = text_decoder_embedding(
 83                input_tokens, current_offset_tensor
 84            )
 85
 86            if isinstance(decoder_inputs["embedded_tokens"], TensorProxy):
 87                decoder_inputs["embedded_tokens"].load_from(
 88                    embedded_tokens, clone=False
 89                )
 90                # mask is mutated at every offset; snapshot it for asynchronous H2D.
 91                decoder_inputs["mask"].load_from(mask, clone=True)
 92                decoder_inputs["offset"].load_from(current_offset_tensor, clone=False)
 93            else:
 94                # CPU/CUDA execution calls decoder_forward directly, so provide
 95                # tensors instead of MLSDK's device-resident input proxies.
 96                decoder_inputs["embedded_tokens"] = embedded_tokens
 97                decoder_inputs["mask"] = mask
 98                decoder_inputs["offset"] = current_offset_tensor
 99
100            outputs = decoder_fn(decoder_inputs)
101            logits = outputs["logits"]
102
103            current_offset += 1
104
105        # transfer logits on device to host
106        logits = logits.cpu()
107        # consider the logits at the last token only
108        logits = logits[:, -1]
109        # apply logit filters
110        for logit_filter in decoding_task.logit_filters:
111            logit_filter.apply(logits, tokens)
112
113        # expand the tokens tensor with the selected next tokens
114        # and judge if all sequences has reached the EOT.
115        tokens, completed = decoding_task.decoder.update(tokens, logits, sum_logprobs)
116
117        if completed or tokens.shape[-1] > model_dims.n_text_ctx:
118            break
119
120    return tokens, sum_logprobs
121
122
123def decode_hypotheses_texts(
124    tokens: torch.Tensor,
125    sum_logprobs: torch.Tensor,
126    decoding_task: DecodingTask,
127) -> list[str]:
128    tokens = tokens.reshape(args.batch_size, decoding_task.n_group, -1)
129    sum_logprobs = sum_logprobs.reshape(args.batch_size, decoding_task.n_group)
130
131    # get the final candidates for each group, and slice between the first sampled token and EOT
132    tokens, sum_logprobs = decoding_task.decoder.finalize(tokens, sum_logprobs)
133
134    def slice_until_eot(t: torch.Tensor) -> torch.Tensor:
135        eot_positions = (t == decoding_task.tokenizer.eot).nonzero()
136        end = int(eot_positions[0, 0]) if eot_positions.numel() else t.shape[0]
137        return t[decoding_task.sample_begin : end]
138
139    tokens = [[slice_until_eot(t) for t in s] for s in tokens]
140
141    # select the top-ranked sample in each group
142    selected = decoding_task.sequence_ranker.rank(tokens, sum_logprobs)
143    tokens = [t[i].tolist() for i, t in zip(selected, tokens)]
144    texts = [decoding_task.tokenizer.decode(t).strip() for t in tokens]
145
146    return texts
147
148
149def transcribe_mels(  # noqa: CFQ002
150    args: argparse.Namespace,
151    audio_encoder_head: AudioEncoderHead,
152    encoder_fn: Callable | CompiledFunction,
153    text_decoder_embedding: TextDecoderEmbedding,
154    text_decoder_transformer: TextDecoderTransformer,
155    decoder_fn: Callable | CompiledFunction,
156    decoder_inputs: dict[str, TensorLike],
157    dataloader: torch.utils.data.DataLoader,
158    encoder_sample: dict[str, torch.Tensor],
159    decoding_task: DecodingTask,
160    context: Context | None,
161) -> tuple[float, list[float]]:
162
163    if context is not None:
164        assert isinstance(decoder_inputs["embedded_tokens"], TensorProxy)
165
166    hypotheses = []
167    references = []
168
169    with Timer() as t:
170
171        total_iterations = len(dataloader)
172        if args.max_iterations != -1:
173            total_iterations = min(total_iterations, args.max_iterations)
174
175        for mels, ref_texts in tqdm(
176            islice(dataloader, total_iterations), total=total_iterations
177        ):
178
179            if mels.ndim == 2:
180                mels = mels.unsqueeze(0)
181
182            decoding_task.decoder.reset()
183            text_decoder_transformer.reset_self_kv_cache(context)
184
185            # whisper.encoder.forward pass
186            embeddings = audio_encoder_head(mels)
187            encoder_sample.update(embeddings=embeddings)
188            cross_kv_cache = encoder_fn(encoder_sample)
189
190            # Compiled encoder outputs are TensorProxy objects and can be passed
191            # directly to the compiled decoder without a device-to-host-to-device
192            # round trip. Native torch backends use the same mapping with tensors.
193            decoder_inputs.update(cross_kv_cache)
194
195            tokens = torch.tensor([decoding_task.initial_tokens]).repeat(
196                args.batch_size, 1
197            )
198            tokens = tokens.repeat_interleave(decoding_task.n_group, dim=0)
199
200            # decode tokens and cumulate
201            tokens, sum_logprobs = decode_loop(
202                text_decoder_embedding,
203                decoder_fn,
204                decoder_inputs,
205                decoding_task,
206                tokens,
207            )
208
209            # decode texts from the inferred results
210            texts = decode_hypotheses_texts(tokens, sum_logprobs, decoding_task)
211            hypotheses.extend(texts)
212            references.extend(ref_texts)
213
214    # Calc Word Error Rate (WER) from the inferred result and the reference data
215    wer = calc_wer(hypotheses, references)
216
217    return wer, [t.time]
218
219
220def run_infer(  # noqa: CFQ001, CFQ002
221    args: argparse.Namespace,
222    audio_encoder_head: AudioEncoderHead,
223    audio_encoder_transformer: AudioEncoderTransformer,
224    text_decoder_embedding: TextDecoderEmbedding,
225    text_decoder_transformer: TextDecoderTransformer,
226    decoding_task: DecodingTask,
227    dataloader: torch.utils.data.DataLoader,
228    context: Context | None,
229    outdir: str,
230) -> tuple[float, list[float]]:
231
232    model_dims = decoding_task.model.dims
233    cross_kv_precompute = CrossKVPrecompute(text_decoder_transformer)
234
235    def encoder_forward(sample_d: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
236        with torch.no_grad():
237            audio_features = audio_encoder_transformer(sample_d["embeddings"])
238            cross_kv_cache = cross_kv_precompute(audio_features)
239        return cross_kv_cache
240
241    def decoder_forward(sample_d: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
242        cross_kv_cache = kv_cache_from_dict(sample_d, model_dims.n_text_layer, "cross")
243
244        with torch.no_grad():
245            logits = text_decoder_transformer(
246                sample_d["embedded_tokens"],
247                sample_d["offset"],
248                sample_d["mask"],
249                cross_kv_cache,
250            )
251
252        return {"logits": logits}
253
254    # prepare sample inputs for compilation of encoder_fn and decoder_fn
255    # The decoder processes one token per invocation in decode_loop().
256    decode_token_length = 1
257    encoder_sample = {
258        "embeddings": torch.randn(
259            args.batch_size, model_dims.n_audio_ctx, model_dims.n_audio_state
260        )
261    }
262    cross_k_cache_dict, cross_v_cache_dict = prepare_cache_sample(args, model_dims)
263
264    decoder_sample = (
265        {
266            "embedded_tokens": torch.randn(
267                args.batch_size,
268                decode_token_length,
269                model_dims.n_text_state,
270                dtype=torch.float32,
271            ),
272            "offset": torch.tensor([1]),
273            "mask": torch.zeros(
274                1, 1, 1, model_dims.n_text_ctx // 2
275            ),  # size of dim=3 is hard-coded
276        }
277        | cross_k_cache_dict
278        | cross_v_cache_dict
279    )
280
281    encoder_fn = encoder_forward
282    decoder_fn = decoder_forward
283    decoder_inputs: dict[str, TensorLike] = dict(decoder_sample)
284    if context is not None:  # for mlsdk
285        # CrossKVPrecompute shares key/value projection parameters with the
286        # decoder. Register the owning decoder module first to avoid registering
287        # the shared tensors twice while compiling the encoder.
288        register_model(context, "whisper_decoder_transformer", text_decoder_transformer)
289
290        # Compile the consumer first so its preferred cross-attention cache
291        # layouts become available to the encoder compilation.
292        print("\n########## compile decoder_fn ###########")
293        decoder_fn = compile_fn(
294            context,
295            decoder_fn,
296            text_decoder_transformer,
297            decoder_sample,
298            outdir=outdir + "_decoder_transformer",
299            model_name="whisper_decoder_transformer",
300            is_train=False,
301            optimize_option=args.optimize_option,
302            preset_options_dir=args.preset_options_dir,
303            enable_cache=not args.disable_compile_cache,
304        )
305
306        decoder_inputs = decoder_fn.allocate_input_proxy()
307        cross_cache_names = set(cross_k_cache_dict) | set(cross_v_cache_dict)
308        for name, tensor in decoder_sample.items():
309            if name not in cross_cache_names:
310                decoder_inputs[name].load_from(tensor, clone=False)
311
312        # Constrain the encoder outputs to the decoder input IOSpecs. This passes
313        # only layout/type metadata to compilation; the encoder and decoder keep
314        # independent allocations, and runtime relocation connects them directly.
315        cross_cache_constraints = {
316            name: decoder_fn.input_specs[name] for name in cross_cache_names
317        }
318
319        print("\n########## compile encoder_fn ##########")
320        encoder_fn = compile_fn(
321            context,
322            encoder_fn,
323            {
324                "whisper_encoder_transformer": audio_encoder_transformer,
325                "whisper_cross_kv_precompute": cross_kv_precompute,
326            },
327            encoder_sample,
328            outdir=outdir + "_encoder_transformer",
329            model_name="whisper_encoder_transformer",
330            is_train=False,
331            optimize_option=args.optimize_option,
332            preset_options_dir=args.preset_options_dir,
333            enable_cache=not args.disable_compile_cache,
334            io_spec_constraints=cross_cache_constraints,
335        )
336
337    return transcribe_mels(
338        args,
339        audio_encoder_head,
340        encoder_fn,
341        text_decoder_embedding,
342        text_decoder_transformer,
343        decoder_fn,
344        decoder_inputs,
345        dataloader,
346        encoder_sample,
347        decoding_task,
348        context,
349    )
350
351
352def main(args: argparse.Namespace) -> None:
353    if args.max_iterations == 0 or args.max_iterations < -1:
354        raise ValueError("max_iterations must be a positive integer or -1")
355    if args.wer_criteria < 0.0 or args.wer_criteria > 1.0:
356        raise ValueError("wer_criteria must be ranged in [0, 1]")
357
358    # Fix seed values for reproducibility
359    set_deterministic_mode(seed=args.seed)
360
361    # Create model obj
362    model = whisper.load_model(args.model, download_root=args.model_dir)
363    model.alignment_heads = model.alignment_heads.to_dense()
364    model.eval()
365    audio_encoder_head = AudioEncoderHead(model.encoder)
366    audio_encoder_transformer = AudioEncoderTransformer(model.encoder)
367    text_decoder_embedding = TextDecoderEmbedding(model.decoder)
368    text_decoder_transformer = TextDecoderTransformer(model.decoder, args.batch_size)
369
370    # Create decoding_task obj to use DecodingTask params in the inference
371    decoding_options = DecodingOptions(language="en", without_timestamps=True)
372    decoding_task = DecodingTask(model, decoding_options)
373
374    # Create Dataset/DataLoader obj
375    infer_dataset = LibriSpeech(
376        dataset_dir=args.dataset_dir,
377        data_name=args.data_name,
378        chunk_length=args.chunk_length,
379        n_mels=model.dims.n_mels,
380        device="cpu",
381    )
382    infer_dataloader = torch.utils.data.DataLoader(
383        infer_dataset, batch_size=args.batch_size, drop_last=True
384    )
385    infer_iterations = len(infer_dataloader)
386    if args.max_iterations != -1:
387        infer_iterations = min(infer_iterations, args.max_iterations)
388
389    # Decide device and outdir from given command line args/options
390    sample_name = "whisper_inference"
391    outdir = decide_outdir(
392        args.backend, example_name=sample_name, basedir=args.out_basedir
393    )
394
395    # Pass device info to the Context obj
396    context = None
397    if args.backend in MLSDK_DEVICES:
398        device = MNDevice(args.backend)
399        context = Context(device)
400        Context.switch_context(context)
401    else:
402        device = args.backend
403
404    with trace_scope(args.trace_filename):
405        # Run inference
406        wer, infer_times = run_infer(
407            args,
408            audio_encoder_head,
409            audio_encoder_transformer,
410            text_decoder_embedding,
411            text_decoder_transformer,
412            decoding_task,
413            infer_dataloader,
414            context,
415            outdir,
416        )
417
418    # Output the inference time
419    output_result_times(
420        eval_times=infer_times,
421        eval_iter=infer_iterations,
422        eval_batch_size=args.batch_size,
423        backend_name=args.backend,
424        sample_name=sample_name,
425        optimize_option=args.optimize_option,
426    )
427
428    print(f"\nResult: WordErrorRate(WER) = {wer * 100:.2f} %")
429    assert wer <= args.wer_criteria, (
430        f"{wer * 100:.2f} % exceeded the "
431        f"specified criteria ({args.wer_criteria * 100:.2f} %)."
432    )
433
434
435if __name__ == "__main__":
436    script_dir = os.path.dirname(__file__)
437    parser = argparse.ArgumentParser(
438        formatter_class=argparse.ArgumentDefaultsHelpFormatter
439    )
440
441    parser.add_argument(
442        "--dataset_dir",
443        type=str,
444        default=f"{script_dir}/librispeech_dataset",
445        help="path to save dataset",
446    )
447    parser.add_argument(
448        "--out_basedir",
449        type=str,
450        default="/tmp",
451        help="basedir to save output artifacts",
452    )
453    parser.add_argument(
454        "--model_dir",
455        type=str,
456        default="/tmp/whisper_models",
457        help="directory in which to download and load Whisper models",
458    )
459    parser.add_argument(
460        "--max_iterations",
461        type=int,
462        default=-1,
463        help="maximum number of DataLoader iterations; -1 processes all samples",
464    )
465    parser.add_argument(
466        "--wer_criteria",
467        type=float,
468        default=0.05,
469        help="acceptable WordErrorRate (WER) criteria",
470    )
471
472    # mlsdk options
473    parser.add_argument(
474        "-b",
475        "--backend",
476        type=str,
477        choices=DEVICES,
478        default="cpu",
479        help="mlsdk options: specify a device as the backend (default: %(default)s)",
480    )
481    parser.add_argument(
482        "--optimize_option",
483        type=str,
484        default="O1",
485        choices=["debug", "O0", "O1", "O2", "O3", "O4"],
486        help="mlsdk options: optimize option for Context.compile() (default %(default)s)",
487    )
488    parser.add_argument(
489        "--preset_options_dir",
490        type=str,
491        default="/opt/pfn/pfcomp/codegen/preset_options",
492        help="mlsdk options: path to preset_options/",
493    )
494    parser.add_argument(
495        "--disable_compile_cache",
496        action="store_true",
497        help="mlsdk options: disable reusing compiled artifacts",
498    )
499    parser.add_argument(
500        "--trace_filename",
501        type=str,
502        default=None,
503        help="mlsdk options: specify a valid filename to enable tracing",
504    )
505
506    apply_toml_defaults(f"{script_dir}/configs.toml", parser)
507
508    args = parser.parse_args()
509
510    main(args)
Listing 8.31 /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference/whisper_utils.py
  1import argparse
  2import os
  3
  4import jiwer
  5import torch
  6import torchaudio
  7import whisper
  8from whisper.model import ModelDimensions
  9from whisper.normalizers import EnglishTextNormalizer
 10
 11
 12# Wrapper class for torchaudio.datasets.LIBRISPEECH class
 13# class definition is originally from
 14# [whisper example](https://github.com/openai/whisper/blob/c0d2f624c09dc18e709e37c2ad90c039a4eb72a2/notebooks/LibriSpeech.ipynb#L56)  # noqa: B950
 15class LibriSpeech(torch.utils.data.Dataset):
 16    def __init__(
 17        self,
 18        dataset_dir: str,
 19        data_name: str = "test-clean",
 20        chunk_length: int = 30,
 21        n_mels: int = 80,
 22        device: str | torch.device = "cpu",
 23    ) -> None:
 24        if not os.path.isdir(dataset_dir):
 25            os.makedirs(dataset_dir)
 26        self.dataset = torchaudio.datasets.LIBRISPEECH(
 27            root=dataset_dir,
 28            url=data_name,
 29            download=True,
 30        )
 31        self.device = device
 32        self.chunk_length = chunk_length
 33        self.n_mels = n_mels
 34        self.SAMPLE_RATE = 16000  # Fixed param
 35
 36    def __len__(self) -> int:
 37        return len(self.dataset)
 38
 39    def __getitem__(self, item: int) -> tuple[torch.Tensor, str]:
 40        audio, sample_rate, text, _, _, _ = self.dataset[item]
 41        assert sample_rate == self.SAMPLE_RATE
 42        audio = whisper.pad_or_trim(
 43            audio.flatten(), length=self.SAMPLE_RATE * self.chunk_length
 44        ).to(self.device)
 45        mel = whisper.log_mel_spectrogram(audio, n_mels=self.n_mels)
 46
 47        return (mel, text)
 48
 49
 50def prepare_cache_sample(
 51    args: argparse.Namespace,
 52    model_dims: ModelDimensions,
 53    dtype: torch.dtype = torch.float32,
 54) -> tuple[
 55    dict[str, torch.Tensor],
 56    dict[str, torch.Tensor],
 57]:
 58
 59    def create_cache_sample(
 60        size: list[int],
 61        n_text_layer: int,
 62        cache_prefix: str,
 63        dtype: torch.dtype = torch.float32,
 64    ) -> dict[str, torch.Tensor]:
 65        return {
 66            f"{cache_prefix}_cache_{i}": torch.zeros(*size, dtype=dtype)
 67            for i in range(n_text_layer)
 68        }
 69
 70    # cache sizes
 71    n_head = model_dims.n_text_head
 72    n_state_per_head = model_dims.n_text_state // model_dims.n_text_head
 73    n_layer = model_dims.n_text_layer
 74    cross_kv_cache_size = [
 75        args.batch_size,
 76        n_head,
 77        model_dims.n_audio_ctx,
 78        n_state_per_head,
 79    ]
 80
 81    cross_k_cache_dict = create_cache_sample(
 82        cross_kv_cache_size, n_layer, "cross_k", dtype
 83    )
 84    cross_v_cache_dict = create_cache_sample(
 85        cross_kv_cache_size, n_layer, "cross_v", dtype
 86    )
 87
 88    return cross_k_cache_dict, cross_v_cache_dict
 89
 90
 91def kv_cache_from_dict(
 92    sample_d: dict[str, torch.Tensor],
 93    n_text_layer: int,
 94    cache_prefix: str,  # self or cross
 95) -> list[tuple[torch.Tensor, torch.Tensor]]:
 96
 97    kv_cache = [
 98        (
 99            sample_d[f"{cache_prefix}_k_cache_{i}"],
100            sample_d[f"{cache_prefix}_v_cache_{i}"],
101        )
102        for i in range(n_text_layer)
103    ]
104
105    return kv_cache
106
107
108def calc_wer(
109    hypotheses: list[str], references: list[str], print_texts: bool = True
110) -> float:
111    normalizer = EnglishTextNormalizer()
112    hypotheses = [normalizer(text) for text in hypotheses]
113    references = [normalizer(text) for text in references]
114
115    if print_texts:
116        print("\n############ Outputs #############")
117        print("\nNormalized references and hypotheses (first 3 samples):")
118        for index, (reference, hypothesis) in enumerate(
119            zip(references[:3], hypotheses[:3])
120        ):
121            print(
122                f"[{index}]\n"
123                f"  Reference : {reference}\n"
124                f"  Hypothesis: {hypothesis}"
125            )
126        print("\n##################################")
127
128    return jiwer.wer(references, hypotheses)
Listing 8.32 /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference/static_whisper.py
  1import types
  2
  3import torch
  4from whisper.model import AudioEncoder, TextDecoder
  5
  6
  7# to whisper.encoder's head modules (conv1(Conv1d) and conv2(Conv1d))
  8class AudioEncoderHead(torch.nn.Module):
  9    def __init__(self, org_encoder: AudioEncoder) -> None:
 10        super().__init__()
 11
 12        self.conv1 = org_encoder.conv1
 13        self.conv2 = org_encoder.conv2
 14
 15    def forward(self, x: torch.Tensor) -> torch.Tensor:
 16        x = torch.nn.functional.gelu(self.conv1(x))
 17        x = torch.nn.functional.gelu(self.conv2(x))
 18        x = x.permute(0, 2, 1)
 19
 20        return x
 21
 22
 23def encoder_self_attn_forward(self, x: torch.Tensor) -> torch.Tensor:
 24    q = self.query(x)
 25    k = self.key(x)
 26    v = self.value(x)
 27
 28    q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
 29    k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
 30    v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
 31
 32    a = torch.nn.functional.scaled_dot_product_attention(q, k, v)
 33    out = a.permute(0, 2, 1, 3).flatten(start_dim=2)
 34
 35    return self.out(out)
 36
 37
 38def encoder_block_forward(self, x: torch.Tensor) -> torch.Tensor:
 39    x = x + self.attn(self.attn_ln(x))
 40    x = x + self.mlp(self.mlp_ln(x))
 41    return x
 42
 43
 44# whisper.encoder's module except to the head
 45class AudioEncoderTransformer(torch.nn.Module):
 46    def __init__(self, org_encoder: AudioEncoder) -> None:
 47        super().__init__()
 48
 49        self.register_buffer("positional_embedding", org_encoder.positional_embedding)
 50        self.blocks = org_encoder.blocks
 51        self.ln_post = org_encoder.ln_post
 52
 53        for block in self.blocks:
 54            # replace ResidualAttentionBlock.forward()
 55            block.forward = types.MethodType(encoder_block_forward, block)
 56            # replace MultiHeadAttention.forward()
 57            block.attn.forward = types.MethodType(encoder_self_attn_forward, block.attn)
 58
 59    def forward(self, x: torch.Tensor) -> torch.Tensor:
 60
 61        assert (
 62            x.shape[1:] == self.positional_embedding.shape
 63        ), f"incorrect audio shape: {x.shape=} vs. {self.positional_embedding.shape=}"
 64
 65        x = (x + self.positional_embedding).to(x.dtype)
 66
 67        for block in self.blocks:
 68            x = block(x)
 69
 70        x = self.ln_post(x)
 71        return x
 72
 73
 74def static_self_attn_forward(
 75    self,
 76    x: torch.Tensor,
 77    mask: torch.Tensor,
 78    kv_cache: tuple[torch.Tensor, torch.Tensor],
 79    offset: torch.Tensor,
 80) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
 81    q = self.query(x)
 82    k = self.key(x)
 83    v = self.value(x)
 84
 85    q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
 86    k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
 87    v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
 88
 89    k_cache = kv_cache[0]
 90    v_cache = kv_cache[1]
 91
 92    # Update the persistent KV cache in place. The cache tensors are registered
 93    # buffers of TextDecoderTransformer, so MLSDK keeps their device allocation
 94    # across decoder invocations.
 95    offset_idx = offset.view(1, 1, 1, 1).expand(*q.shape)
 96    k_cache.scatter_(dim=2, index=offset_idx, src=k)
 97    v_cache.scatter_(dim=2, index=offset_idx, src=v)
 98
 99    a = torch.nn.functional.scaled_dot_product_attention(
100        q,
101        k_cache,
102        v_cache,
103        attn_mask=mask,
104        dropout_p=0.0,
105        is_causal=False,
106    )
107    out = a.permute(0, 2, 1, 3).flatten(start_dim=2)
108
109    return self.out(out)
110
111
112def static_cross_attn_forward(
113    self,
114    x: torch.Tensor,
115    kv_cache: tuple[torch.Tensor, torch.Tensor],
116) -> torch.Tensor:
117    q = self.query(x)
118    q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
119    # view() and permute() have been done in CrossKVPrecompute.
120    k, v = kv_cache
121
122    a = torch.nn.functional.scaled_dot_product_attention(q, k, v)
123    out = a.permute(0, 2, 1, 3).flatten(start_dim=2)
124
125    return self.out(out)
126
127
128def static_block_forward(
129    self,
130    x: torch.Tensor,
131    mask: torch.Tensor | None = None,
132    kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
133    offset: torch.Tensor | None = None,
134    cross_kv_cache: (
135        tuple[torch.Tensor, torch.Tensor] | None
136    ) = None,  # should be precomputed
137) -> torch.Tensor:
138    # apply self_attention
139    attn_out = self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache, offset=offset)
140    x = x + attn_out
141
142    # apply cross_attention
143    if self.cross_attn:
144        x = x + self.cross_attn(self.cross_attn_ln(x), kv_cache=cross_kv_cache)
145
146    # apply mlp
147    x = x + self.mlp(self.mlp_ln(x))
148
149    return x
150
151
152# to calculation of the embeddings of the input tokens
153class TextDecoderEmbedding(torch.nn.Module):
154    def __init__(self, org_decoder: TextDecoder) -> None:
155        super().__init__()
156        self.token_embedding = org_decoder.token_embedding
157        self.positional_embedding = org_decoder.positional_embedding
158
159    def forward(self, x: torch.Tensor, offset: torch.Tensor) -> torch.Tensor:
160        pos_emb = self.positional_embedding[offset[0]].view(
161            1, 1, -1
162        )  # to broadcast explictly
163        x = self.token_embedding(x) + pos_emb
164
165        return x
166
167
168# whisper.decoder's module except to embeddings
169class TextDecoderTransformer(torch.nn.Module):
170    def __init__(self, org_decoder: TextDecoder, batch_size: int) -> None:
171        super().__init__()
172
173        self.blocks = org_decoder.blocks
174        self.ln = org_decoder.ln
175        self.token_embedding = org_decoder.token_embedding
176
177        # TextDecoder does not expose n_ctx; its positional embedding has one
178        # row per text-context position.
179        cache_length = org_decoder.positional_embedding.shape[0] // 2
180        n_head = self.blocks[0].attn.n_head
181        n_state_per_head = self.token_embedding.embedding_dim // n_head
182        cache_shape = (batch_size, n_head, cache_length, n_state_per_head)
183        for i in range(len(self.blocks)):
184            self.register_buffer(f"self_k_cache_{i}", torch.zeros(cache_shape))
185            self.register_buffer(f"self_v_cache_{i}", torch.zeros(cache_shape))
186
187        for block in self.blocks:
188            # replace forward() of ResidualAttentionBlock
189            block.forward = types.MethodType(static_block_forward, block)
190
191            # replace forward() of self_attention
192            block.attn.forward = types.MethodType(static_self_attn_forward, block.attn)
193
194            # replace forward() of cross_attention
195            block.cross_attn.forward = types.MethodType(
196                static_cross_attn_forward, block.cross_attn
197            )
198
199    def self_kv_cache(self) -> list[tuple[torch.Tensor, torch.Tensor]]:
200        return [
201            (
202                getattr(self, f"self_k_cache_{i}"),
203                getattr(self, f"self_v_cache_{i}"),
204            )
205            for i in range(len(self.blocks))
206        ]
207
208    def reset_self_kv_cache(self, context=None) -> None:
209        """Clear self-attention state before decoding the next audio sample."""
210        for k_cache, v_cache in self.self_kv_cache():
211            k_cache.zero_()
212            v_cache.zero_()
213            if context is not None:
214                context.get_registered_value_proxy(k_cache).load_from(
215                    k_cache, clone=False
216                )
217                context.get_registered_value_proxy(v_cache).load_from(
218                    v_cache, clone=False
219                )
220
221    def forward(
222        self,
223        x: torch.Tensor,  # embedded_tokens
224        offset: torch.Tensor,  # offset for kv_cache
225        mask: torch.Tensor,  # mask for kv_cache
226        cross_kv_cache: list[tuple[torch.Tensor, torch.Tensor]],  # len(list) == n_layer
227    ) -> torch.Tensor:
228
229        x = x.to(cross_kv_cache[0][0].dtype)
230
231        self_kv_cache = self.self_kv_cache()
232        for i, block in enumerate(self.blocks):
233            self_kv = self_kv_cache[i]
234            cross_kv = cross_kv_cache[i]
235            x = block(
236                x,
237                mask=mask,
238                kv_cache=self_kv,
239                offset=offset,
240                cross_kv_cache=cross_kv,
241            )
242
243        x = self.ln(x)
244        logits = (
245            x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)
246        ).float()
247
248        return logits
249
250
251class CrossKVPrecompute(torch.nn.Module):
252    """Precompute decoder cross-attention K/V caches from audio features."""
253
254    def __init__(self, text_decoder_transformer: TextDecoderTransformer) -> None:
255        super().__init__()
256
257        # These projection modules are shared with TextDecoderTransformer.
258        # Register TextDecoderTransformer with the MLSDK Context before compiling
259        # the encoder so that their parameters are registered only once.
260        self.keys = torch.nn.ModuleList(
261            [block.cross_attn.key for block in text_decoder_transformer.blocks]
262        )
263        self.values = torch.nn.ModuleList(
264            [block.cross_attn.value for block in text_decoder_transformer.blocks]
265        )
266        self.n_heads = [
267            block.cross_attn.n_head for block in text_decoder_transformer.blocks
268        ]
269
270    def forward(self, audio_features: torch.Tensor) -> dict[str, torch.Tensor]:
271        cross_kv_cache = {}
272
273        for i, (key, value, n_head) in enumerate(
274            zip(self.keys, self.values, self.n_heads)
275        ):
276            k = key(audio_features)
277            v = value(audio_features)
278
279            k = k.view(k.shape[0], k.shape[1], n_head, -1).permute(0, 2, 1, 3)
280            v = v.view(v.shape[0], v.shape[1], n_head, -1).permute(0, 2, 1, 3)
281
282            cross_kv_cache[f"cross_k_cache_{i}"] = k
283            cross_kv_cache[f"cross_v_cache_{i}"] = v
284
285        return cross_kv_cache
Listing 8.33 /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference/utility.py
  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("")
Listing 8.34 /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference/requirements.txt
1openai-whisper
2soundfile
3jiwer
4pandas
5torchaudio==2.9.0
6torchcodec==0.9.0
Listing 8.35 /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference/requirements_cpu.txt
1torchaudio==2.9.0
2torchcodec==0.9.0
Listing 8.36 /opt/pfn/pfcomp/codegen/MLSDK/examples/whisper_inference/configs.toml
 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