def benchmark_all() -> pd.DataFrame:
save_path = Path("./benchmark.csv")
df = pd.DataFrame()
if save_path.exists():
df = pd.read_csv(str(save_path), index_col=0)
index = []
records = []
bar = tqdm(models)
try:
for key in bar:
if key not in df.index:
try:
model = AutoModel.from_pretrained(key)
tr = AutoTransform.from_name(key)
batch_size = 64
if key in batch_sizes:
batch_size = batch_sizes[key]
bar.set_description(
f"{key}, size={tr.transforms[0].size}, batch_size={batch_size}"
)
top1, top5, time = benchmark(model.to(device), tr,
batch_size)
index.append(key)
data = {
"top1": top1,
"top5": top5,
"time": time,
"batch_size": batch_size,
}
pprint(data)
records.append(data)
except KeyError:
continue
except Exception as e:
print(e)
pass
if len(records) > 0:
new_df = pd.DataFrame.from_records(records, index=index)
if df is not None:
df = pd.concat([df, new_df])
else:
df = new_df
df.to_csv("./benchmark.csv")
mk = df.sort_values("top1", ascending=False).to_markdown()
with open("./benchmark.md", "w") as f:
f.write(mk)
return df
def vit_clone(key: str):
src = timm.create_model(key, pretrained='True')
dst = AutoModel.from_name(key)
dst.embedding.positions.data.copy_(src.pos_embed.data.squeeze(0))
dst.embedding.cls_token.data.copy_(src.cls_token.data)
cfg = AutoConfig.from_name(key)
return clone_model(src, dst,
torch.randn((1, 3, cfg.input_size, cfg.input_size)))
def benchmark_all() -> pd.DataFrame:
save_path = Path('./benchmark.csv')
df = pd.DataFrame()
if save_path.exists():
df = pd.read_csv(str(save_path), index_col=0)
index = []
records = []
bar = tqdm(models)
try:
for key in bar:
if key not in df.index:
try:
model = AutoModel.from_pretrained(key)
cfg = AutoConfig.from_name(key)
tr = cfg.transform
batch_size = 64
# if key in batch_sizes:
# batch_size = batch_sizes[key]
bar.set_description(
f'{key}, size={cfg.input_size}, batch_size={batch_size}'
)
top1, top5, time = benchmark(model.to(device), tr,
batch_size)
index.append(key)
data = {
'top1': top1,
'top5': top5,
'time': time,
'batch_size': batch_size
}
pprint(data)
records.append(data)
except KeyError:
continue
except Exception as e:
print(e)
pass
new_df = pd.DataFrame.from_records(records, index=index)
if df is not None:
df = pd.concat([df, new_df])
else:
df = new_df
df.to_csv('./benchmark.csv')
print(df)
def vit_clone(key: str):
src = timm.create_model(key, pretrained="True")
dst = AutoModel.from_name(key)
cfg = AutoTransform.from_name(key)
dst = clone_model(
src,
dst,
torch.randn((1, 3, cfg.input_size, cfg.input_size)),
dest_skip=[ViTTokens],
)
dst.embedding.positions.data.copy_(src.pos_embed.data.squeeze(0))
dst.embedding.tokens.cls.data.copy_(src.cls_token.data)
return dst
def deit_clone(key: str):
k_split = key.split('_')
hub_key = "_".join(k_split[:2]) + '_distilled_' + "_".join(k_split[2:])
src = torch.hub.load('facebookresearch/deit:main',
hub_key, pretrained=True)
dst = AutoModel.from_name(key)
cfg = AutoConfig.from_name(f"vit_{'_'.join(key.split('_')[1:])}")
dst = clone_model(src, dst, torch.randn(
(1, 3, cfg.input_size, cfg.input_size)), dest_skip=[DeiTTokens])
dst.embedding.positions.data.copy_(src.pos_embed.data.squeeze(0))
dst.embedding.tokens.cls.data.copy_(src.cls_token.data)
dst.embedding.tokens.dist.data.copy_(src.dist_token.data)
return dst
with open("pretrained_models.txt", "w") as f:
f.write(",".join(list(zoo_source.keys())))
if args.o is not None:
save_dir = args.o
save_dir.mkdir(exist_ok=True)
storages = {"local": LocalStorage, "hf": HuggingFaceStorage}
storage = storages[args.storage]()
if args.storage == "local":
logging.info(f"Store root={storage.root}")
override = True
bar = tqdm(zoo_source.items())
uploading_bar = tqdm()
for key, src_def in bar:
bar.set_description(key)
if src_def is None:
# it means I was lazy and I meant to use timm
src_def = partial(timm.create_model, key, pretrained=True)
if key not in storage or override:
if type(src_def) is tuple:
# I have a custom clone func -> not the most elegant way, but it works!
clone_func, flag = src_def
cloned = clone_func(key)
else:
src, dst = src_def(), AutoModel.from_name(key)
cloned = clone_model(src, dst)
storage.put(key, cloned)