zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-06-24 17:31:18 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Benchmark all models if user allows.

Browse Source
This commit is contained in:
adhishthite 2023-12-07 00:07:42 +05:30
parent 0bf5d0e3bc
commit 9cf82a0d43
2 changed files with 58 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@ -127,3 +127,5 @@ dmypy.json
# Pyre type checker
.pyre/
.idea/
.vscode/

99
whisper/benchmark.py
View File

@ -1,5 +1,6 @@
# Copyright © 2023 Apple Inc.
import sys
import time
import mlx.core as mx
@ -48,46 +49,58 @@ def everything():
if __name__ == "__main__":
feat_time = timer(feats)
print(f"Feature time {feat_time:.3f}")
mels = feats()[None]
tokens = mx.array(
[
50364,
1396,
264,
665,
5133,
23109,
25462,
264,
6582,
293,
750,
632,
42841,
292,
370,
938,
294,
4054,
293,
12653,
356,
50620,
50620,
23563,
322,
3312,
13,
50680,
],
mx.int32,
)[None]
model = load_models.load_model("tiny")
model_forward_time = timer(model_forward, model, mels, tokens)
print(f"Model forward time {model_forward_time:.3f}")
decode_time = timer(decode, model, mels)
print(f"Decode time {decode_time:.3f}")
everything_time = timer(everything)
print(f"Everything time {everything_time:.3f}")
# get command line arguments without 3rd party libraries
# the command line argument to benchmark all models is "all"
models = ["tiny"]
if len(sys.argv) > 1:
if sys.argv[1] == "--all":
models = ["tiny", "small", "medium", "large"]
for model_name in models:
feat_time = timer(feats)
print(f"\nModel: {model_name.upper()}")
print(f"\nFeature time {feat_time:.3f}")
mels = feats()[None]
tokens = mx.array(
[
50364,
1396,
264,
665,
5133,
23109,
25462,
264,
6582,
293,
750,
632,
42841,
292,
370,
938,
294,
4054,
293,
12653,
356,
50620,
50620,
23563,
322,
3312,
13,
50680,
],
mx.int32,
)[None]
model = load_models.load_model(f"{model_name}")
model_forward_time = timer(model_forward, model, mels, tokens)
print(f"Model forward time {model_forward_time:.3f}")
decode_time = timer(decode, model, mels)
print(f"Decode time {decode_time:.3f}")
everything_time = timer(everything)
print(f"Everything time {everything_time:.3f}")
print(f"\n{'-----' * 10}\n")