Skip to content

TimmWrapper

PyTorch

Helper class to enable loading timm models to be used with the transformers library and its autoclasses.

>>> import torch
>>> from PIL import Image
>>> from urllib.request import urlopen
>>> from transformers import AutoModelForImageClassification, AutoImageProcessor
>>> # Load image
>>> image = Image.open(urlopen(
... 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
... ))
>>> # Load model and image processor
>>> checkpoint = "timm/resnet50.a1_in1k"
>>> image_processor = AutoImageProcessor.from_pretrained(checkpoint)
>>> model = AutoModelForImageClassification.from_pretrained(checkpoint).eval()
>>> # Preprocess image
>>> inputs = image_processor(image)
>>> # Forward pass
>>> with torch.no_grad():
... logits = model(**inputs).logits
>>> # Get top 5 predictions
>>> top5_probabilities, top5_class_indices = torch.topk(logits.softmax(dim=1) * 100, k=5)

A list of official Hugging Face and community (indicated by 🌎) resources to help you get started with TimmWrapper.

[[autodoc]] TimmWrapperConfig

[[autodoc]] TimmWrapperImageProcessor - preprocess

[[autodoc]] TimmWrapperModel - forward

[[autodoc]] TimmWrapperForImageClassification - forward