Skip to content

Instantly share code, notes, and snippets.

@olokobayusuf

olokobayusuf/mobilenet.py Secret

Last active September 28, 2021 16:44
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 olokobayusuf/42ca58275037ab6627cd3dd6a5b539e4 to your computer and use it in GitHub Desktop.
Save olokobayusuf/42ca58275037ab6627cd3dd6a5b539e4 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
mobilenet.py
# Load the pre-trained MobileNet v3 model
model = torchvision.models.mobilenet_v3_small(pretrained=True).eval()
# Perform inference
input = torch.randn(1, 3, 224, 224)
logits = model(input) # shape: (1,1000)
# Get the output label
result = logits.argmax(dim=1).item()
labels = ["cat", "dog", ...]
result_label = labels[result]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment