File size: 910 Bytes
8a6022e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
from model import DINOv2FeatureExtractor
import torch
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
MODEL_CHECKPOINT_PATH = './weights/best_model_95.6.torch'
model = DINOv2FeatureExtractor(
model_type="vit_base_patch14_reg4_dinov2.lvd142m",
num_of_layers_to_unfreeze=0,
desc_dim=768,
aggregator_type="SALAD",
)
print('loading model ... ')
model_state_dict = torch.load(MODEL_CHECKPOINT_PATH, map_location=DEVICE)
model.load_state_dict(model_state_dict)
model = model.to(DEVICE)
model.eval()
print('loaded ....')
# Move to device
model.to(DEVICE)
# Print some info about model weights
num_params = sum(p.numel() for p in model.parameters())
num_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model total parameters: {num_params:,}")
print(f"Model trainable parameters: {num_trainable:,}")
print(model.aggregator_type)
|