distillation_bench.py

import time
import torch
import open_clip

from data import get_data
from params import parse_args


# Args
args = parse_args()
dev = "cuda"

# Model
model_l, _, preprocess = open_clip.create_model_and_transforms('ViT-L-14', pretrained='laion400m_e32')
model_l = model_l.to(dev)

images = torch.randn((256, 3, 224, 224)).to(dev)
texts = torch.ones((256, 77), dtype=int).to(dev)
for i in range(5):
    t0 = time.time()
    with torch.no_grad():
        _, _, _ = model_l(images, texts)
    tf = time.time()
    print(images.shape[0]/(tf-t0))