Skip to content

Instantly share code, notes, and snippets.

@AntreasAntoniou

AntreasAntoniou/vit_tiny_84x84.py

Last active March 6, 2023 15:56
Show Gist options
  • Save AntreasAntoniou/99203963e622dd508f41ca0efc6a8290 to your computer and use it in GitHub Desktop.
Save AntreasAntoniou/99203963e622dd508f41ca0efc6a8290 to your computer and use it in GitHub Desktop.
Download ZIP
ViT tiny demo
Raw
vit_tiny_84x84.py
from rich import print
import timm
import torch
import tqdm
model_names = ["vit_base_patch32_224_clip_laion2b", "vit_tiny_patch16_224_in21k"]
model = timm.create_model(
"vit_base_patch32_224_clip_laion2b", pretrained=True, in_chans=3, img_size=84
).to(torch.cuda.current_device())
dummy_input = torch.randn(16 * 180, 3, 84, 84).to(torch.cuda.current_device())
out = model.forward_features(dummy_input)
print(out.shape)
with tqdm.tqdm(10000) as t:
for i in range(10000):
out = model.forward_features(dummy_input)
loss = out.mean()
loss.backward()
t.update(1)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment