Views
No views yet
.pth to .safetensors for safer and faster loading.SimCLR class definition from the accompanying files.pip install torch torchvision safetensors huggingface_hub1import torch
2import json
3from safetensors.torch import load_file
4from huggingface_hub import hf_hub_download
5from models.simclr import SimCLR # Ensure your model code is in the path
6
7# 1. Download files
8repo_id = "homeboi/luthra_simclr_im1k_r50"
9config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
10weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")
11
12# 2. Setup Configuration
13with open(config_path, "r") as f:
14 config = json.load(f)
15
16# 3. Initialize Encoder
17if config["encoder_type"] == 'resnet50':
18 import torchvision.models as models
19 encoder = models.resnet50(pretrained=False)
20elif config["encoder_type"] == 'vit_b':
21 PATCH_SIZE = config["patch_size"]
22 IMAGE_SIZE = config["image_size"]
23 HIDDEN_DIM = config["token_hidden_dim"]
24 MLP_DIM = config["mlp_dim"]
25 STRIDE = config["stride"]
26 encoder = models.VisionTransformer(
27 patch_size=PATCH_SIZE,
28 image_size=IMAGE_SIZE,
29 hidden_dim=HIDDEN_DIM,
30 mlp_dim=MLP_DIM,
31 num_layers=12,
32 num_heads=12,
33 )
34
35# 4. Initialize SimCLR
36model = SimCLR(
37 model=encoder,
38 dataset=config["dataset"],
39 width_multiplier=config["width_multiplier"],
40 hidden_dim=config["hidden_dim"],
41 projection_dim=config["projection_dim"],
42 image_size=config["image_size"],
43 patch_size=config["patch_size"],
44 stride=config["stride"],
45 token_hidden_dim=config["token_hidden_dim"],
46 mlp_dim=config["mlp_dim"]
47)
48
49# 5. Load Safetensors
50state_dict = load_file(weights_path)
51model.load_state_dict(state_dict)
52model.eval()
53
54print("Model loaded successfully from Safetensors!")