zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-01 12:49:50 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

whisper default in fp16

Browse Source
This commit is contained in:
Awni Hannun
2023-12-12 07:37:35 -08:00
parent 13f1142eaa
commit 6e723a015a
6 changed files with 50 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
whisper/benchmark.py
View File

@@ -57,12 +57,13 @@ if __name__ == "__main__":
if sys.argv[1] == "--all":
models = ["tiny", "small", "medium", "large"]
feat_time = timer(feats)
print(f"\nFeature time {feat_time:.3f}")
mels = feats()[None].astype(mx.float16)
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,
@@ -96,7 +97,7 @@ if __name__ == "__main__":
],
mx.int32,
)[None]
model = load_models.load_model(f"{model_name}")
model = load_models.load_model(f"{model_name}", dtype=mx.float16)
model_forward_time = timer(model_forward, model, mels, tokens)
print(f"Model forward time {model_forward_time:.3f}")
decode_time = timer(decode, model, mels)