zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-06-25 01:41:19 +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 type checker
.pyre/ .pyre/
.idea/
.vscode/

17
whisper/benchmark.py
View File

@ -1,5 +1,6 @@
# Copyright © 2023 Apple Inc. # Copyright © 2023 Apple Inc.
import sys
import time import time
import mlx.core as mx import mlx.core as mx
@ -48,8 +49,19 @@ def everything():
if __name__ == "__main__": if __name__ == "__main__":
# 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) feat_time = timer(feats)
print(f"Feature time {feat_time:.3f}")
print(f"\nModel: {model_name.upper()}")
print(f"\nFeature time {feat_time:.3f}")
mels = feats()[None] mels = feats()[None]
tokens = mx.array( tokens = mx.array(
[ [
@ -84,10 +96,11 @@ if __name__ == "__main__":
], ],
mx.int32, mx.int32,
)[None] )[None]
model = load_models.load_model("tiny") model = load_models.load_model(f"{model_name}")
model_forward_time = timer(model_forward, model, mels, tokens) model_forward_time = timer(model_forward, model, mels, tokens)
print(f"Model forward time {model_forward_time:.3f}") print(f"Model forward time {model_forward_time:.3f}")
decode_time = timer(decode, model, mels) decode_time = timer(decode, model, mels)
print(f"Decode time {decode_time:.3f}") print(f"Decode time {decode_time:.3f}")
everything_time = timer(everything) everything_time = timer(everything)
print(f"Everything time {everything_time:.3f}") print(f"Everything time {everything_time:.3f}")
print(f"\n{'-----' * 10}\n")