|
|
|
from huggingface_hub import hf_hub_download |
|
from vision_transformer import vit_large_patch16_224_in21k |
|
import torch |
|
import numpy as np |
|
|
|
REPO_ID = "ethz-mtc/aesthetics_vit" |
|
FILENAME="pytorch_model.bin" |
|
|
|
path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME, cache_dir=".models") |
|
print(path) |
|
REPO_ID = "ethz-mtc/shot_scale_classifier-resnet50" |
|
path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME, cache_dir=".models") |
|
print(path) |
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
|
|
|
model = vit_large_patch16_224_in21k() |
|
model.reset_classifier(num_classes=1) |
|
model.load_state_dict(torch.load(path, map_location=device)) |
|
|
|
print( |
|
f"Model has {sum(np.prod(p.shape) for p in model.parameters()):,} parameters." |
|
) |
