Views
No views yet
"checkpoint_{i}.pth", where i the the number of traineng sample the model has been trained on.1!git clone https://github.com/soniajoseph/ViT-Prisma
2!cd ViT-Prisma
3!pip install -e .1from huggingface_hub import hf_hub_download
2import torch
3
4REPO_ID = "IamYash/dSprites-tiny-AttentionOnly"
5FILENAME = "model_0.pth"
6
7checkpoint = torch.load(
8hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
9)1from vit_prisma.models.base_vit import BaseViT
2from vit_prisma.configs.DSpritesConfig import GlobalConfig
3from vit_prisma.utils.wandb_utils import update_dataclass_from_dict
4
5config = GlobalConfig()
6print(config)
7update_dict = {
8 'transformer':{
9 'attention_only': True,
10 'hidden_dim': 512,
11 'num_heads': 8,
12 'num_layers': 1
13 }
14}
15update_dataclass_from_dict(config, update_dict)
16
17model = BaseViT(config)
18
19model.load_state_dict(checkpoint['model_state_dict'])