Skip to content

Instantly share code, notes, and snippets.

@maziyarpanahi

maziyarpanahi/hf-vit-pipeline-base.py

Created August 14, 2022 16:06
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save maziyarpanahi/f45ac2767e6e69cb17e97961d47be4f1 to your computer and use it in GitHub Desktop.
Save maziyarpanahi/f45ac2767e6e69cb17e97961d47be4f1 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
hf-vit-pipeline-base.py
from transformers import ViTFeatureExtractor, ViTForImageClassification
from transformers import pipeline
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
pipe = pipeline("image-classification", model=model, feature_extractor=feature_extractor, device=-1)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment