Skip to content

Instantly share code, notes, and snippets.

@maziyarpanahi

maziyarpanahi/hf-vit-pipeline-bench.py

Created August 14, 2022 16:16
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save maziyarpanahi/6f684f2fa23853a7565a4e61ce122adf to your computer and use it in GitHub Desktop.
Save maziyarpanahi/6f684f2fa23853a7565a4e61ce122adf to your computer and use it in GitHub Desktop.
Download ZIP
Raw
hf-vit-pipeline-bench.py
from transformers import pipeline
pipe = pipeline("image-classification", model=model, feature_extractor=feature_extractor, device=-1)
for batch_size in [1, 8, 32, 64, 128]:
print("-" * 30)
print(f"Streaming batch_size={batch_size}")
for out in tqdm(pipe(dataset, batch_size=batch_size), total=len(dataset)):
pass
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment